diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 582961f8..7d3160ea 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -77,7 +77,7 @@ jobs: run: docker build -t ade-bench-base -f docker/base/Dockerfile.duckdb-dbt . - name: Run benchmark - run: uv run ade run all --agent ${{ matrix.agent }} --db duckdb --project-type dbt --no-diffs --n-concurrent-trials 6 --no-rebuild + run: uv run ade run all --agent ${{ matrix.agent }} --db duckdb --project-type dbt --no-diffs --n-concurrent-trials 6 --no-rebuild --plugin-set none env: USE_DYNAMIC_LOGGING: "FALSE" DEFAULT_TEST_TIMEOUT_SEC: "120" diff --git a/CLAUDE.md b/CLAUDE.md index f5aac999..f5e1585d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -50,12 +50,13 @@ uv run scripts_python/run_harness.py --agent sage --task-ids task1 task2 ``` ### Key Parameters: -- `--agent`: Agent type (sage, terminus, etc.) +- `--agent`: Agent type (sage, claude, codex, gemini, etc.) - `--model`: LLM model for AI agents - `--dataset-config`: YAML file defining task collection - `--n-concurrent-trials`: Parallel execution (default: 4) - `--no-rebuild`: Skip Docker rebuilds - `--cleanup`: Remove Docker resources after run +- `--plugin-set`: Plugin set names from `experiment_sets/plugin-sets.yaml` (space-separated). Controls skills, MCP servers, and allowed tools. Defaults to `none` (no plugins). ## Development Workflow diff --git a/README.md b/README.md index 1b920e99..07333812 100644 --- a/README.md +++ b/README.md @@ -112,7 +112,7 @@ ade run all --db duckdb --project-type dbt --agent claude ### 8. Go beyond - Use the dbt Fusion engine instead of dbt Core with `--project-type dbt-fusion` ([set up Snowflake](#snowflake-setup) first) -- Enable the dbt MCP server with the `--use-mcp` flag (requires Snowflake, see [MCP](#enabling-the-mcp-server) section) +- Enable skills, MCP servers, or both with the `--plugin-set` flag (see [Plugin Sets](#plugin-sets) section) - [Contribute additional tasks or datasets](/docs/CONTRIBUTING.md) --- @@ -159,7 +159,7 @@ ade run \ --seed \ # Optional; flag for creating solution seed CSVs !! DESTRUCTIVE !! RUN WITH CAUTION !! SEE BELOW !! --no-diffs \ # Optional; disables taking snapshots of diffs for faster performance. --persist \ # Optional; keeps the container alive after the trial is over or is aborted. - --use-mcp \ # Optional; creates an dbt MCP server for the agent. Note: Not all agents and databases are supported. + --plugin-set none \ # Optional; plugin set name(s) from experiment_sets/plugin-sets.yaml. Defaults to 'none'. --tasks-dir /absolute/path/to/tasks \ # Optional; path to an external tasks directory. Defaults to 'tasks' in the current directory. ``` @@ -397,14 +397,26 @@ gemini --output-format json --yolo --prompt {task_prompt} --model {model-id} Configuration files for each agent are found in the `/shared/config` directory. You can use `CLAUDE.md` to configure Claude Code, `AGENTS.md` to configure Codex, and `GEMINI.md` to configure Gemini. -### Enabling the MCP server +### Plugin sets -If run with the flag `--use-mcp`, ADE-bench will create a dbt MCP server that the agent is allowed to use. The following databases and agents are supported: +Plugin sets are declarative configurations of skills, MCP servers, and allowed tools that can be applied to benchmark runs. They are defined in `experiment_sets/plugin-sets.yaml`. -- Databases: `snowflake` (duckdb doesn't support multiple simultaneous connections) -- Agents: `claude`, `codex`, `gemini` +Use the `--plugin-set` flag to select one or more plugin sets: -Because the server runs locally, it only has access to the [CLI tools](https://github.com/dbt-labs/dbt-mcp#tools). The others are disabled, because they require access to the dbt platform. +```bash +ade run all --db duckdb --project-type dbt --agent claude --plugin-set all-dbt-skills +ade run all --db snowflake --project-type dbt --agent claude --plugin-set dbt-mcp +ade run all --db duckdb --project-type dbt --agent claude --plugin-set none # baseline (default) +``` + +Available plugin sets include: +- `none` (default): No skills or MCP servers. Baseline configuration. +- `all-dbt-skills`: Installs all skills from `dbt-labs/dbt-agent-skills`. +- `dbt-for-ae`: Installs only the `using-dbt-for-analytics-engineering` skill. +- `dbt-mcp`: Configures the dbt MCP server (Snowflake only — DuckDB doesn't support multiple simultaneous connections). Only [CLI tools](https://github.com/dbt-labs/dbt-mcp#tools) are enabled. +- `dbt-skills-mcp`: Both dbt skills and the dbt MCP server. + +See `experiment_sets/plugin-sets.yaml` for the full configuration of each set. ### The Sage agent diff --git a/ade_bench/agents/base_agent.py b/ade_bench/agents/base_agent.py index a31b8c6a..8f543387 100644 --- a/ade_bench/agents/base_agent.py +++ b/ade_bench/agents/base_agent.py @@ -75,6 +75,21 @@ def format_agent_log(self, log_path: Path) -> str | None: """ return None + def extract_tools_used(self, log_path: Path) -> list[str] | None: + """ + Extract deduplicated list of tool names from the agent's log file. + + This method can be overridden by subclasses to provide agent-specific + tool extraction. The default implementation returns None. + + Args: + log_path: Path to the raw agent log file + + Returns: + Sorted list of unique tool names, or None if not available + """ + return None + @abstractmethod def perform_task( self, diff --git a/ade_bench/agents/installed_agents/abstract_installed_agent.py b/ade_bench/agents/installed_agents/abstract_installed_agent.py index 04a25a69..0c08e9b8 100644 --- a/ade_bench/agents/installed_agents/abstract_installed_agent.py +++ b/ade_bench/agents/installed_agents/abstract_installed_agent.py @@ -18,7 +18,9 @@ from ade_bench.agents.agent_name import AgentName from ade_bench.agents.base_agent import AgentResult, BaseAgent from ade_bench.harness_models import TerminalCommand, FailureMode +from ade_bench.harness_models import McpServerConfig from ade_bench.terminal.tmux_session import TmuxSession +from ade_bench.terminal.docker_compose_manager import DockerComposeManager from ade_bench.utils.logger import log_harness_info, logger from ade_bench.config import config @@ -26,11 +28,18 @@ class AbstractInstalledAgent(BaseAgent, ABC): NAME = AgentName.ABSTRACT_INSTALLED - def __init__(self, use_mcp: bool = False, model_name: str | None = None, **kwargs): + def __init__( + self, + model_name: str | None = None, + allowed_tools: list[str] | None = None, + mcp_servers: dict[str, McpServerConfig] | None = None, + **kwargs, + ): super().__init__(**kwargs) self._variant_config = {} - self._use_mcp = use_mcp self._model_name = model_name + self._allowed_tools = allowed_tools or [] + self._mcp_servers = mcp_servers or {} @property @abstractmethod @@ -61,6 +70,95 @@ def _run_agent_commands(self, task_prompt: str) -> list[TerminalCommand]: def _create_env_setup_file(self) -> str: return "\n".join([f"export {key}='{value}'" for key, value in self._env.items()]) + def _get_dbt_dynamic_env(self, session: TmuxSession, task_name: str | None) -> dict[str, str]: + """Get dynamic environment variables for dbt MCP server.""" + env_vars = {} + + # DBT_PROJECT_DIR is the container app directory + env_vars["DBT_PROJECT_DIR"] = str(DockerComposeManager.CONTAINER_APP_DIR) + + # Get the dbt path from the container + result = session.container.exec_run( + ["sh", "-c", "which dbt"], workdir=str(DockerComposeManager.CONTAINER_APP_DIR) + ) + if result.exit_code == 0: + dbt_path = result.output.decode("utf-8").strip() + if dbt_path: + env_vars["DBT_PATH"] = dbt_path + log_harness_info(logger, task_name, "agent", f"Found dbt at: {dbt_path}") + else: + logger.warning("[MCP] dbt not found in PATH, MCP server may not work correctly") + + # Enable the dbt CLI in dbt-mcp + env_vars["DISABLE_DBT_CLI"] = "false" + + return env_vars + + def _configure_mcp_servers(self, session: TmuxSession, task_name: str | None) -> None: + """Configure MCP servers after agent installation.""" + agent_cli = self.NAME.value # e.g., "claude", "gemini" + + for server_name, mcp_config in self._mcp_servers.items(): + log_harness_info( + logger, task_name, "agent", f"Configuring MCP server '{server_name}'..." + ) + + # Start with static env vars from config + env_vars = dict(mcp_config.env) + + # For dbt MCP server, add dynamic environment variables + # Check server name or if dbt-mcp appears in any of the args + is_dbt_mcp = server_name == "dbt" or any("dbt-mcp" in arg for arg in mcp_config.args) + if is_dbt_mcp: + dynamic_env = self._get_dbt_dynamic_env(session, task_name) + # Merge dynamic vars (don't override static config) + for key, value in dynamic_env.items(): + if key not in env_vars: + env_vars[key] = value + + # Write env file if we have any env vars + env_file_path = None + if env_vars: + env_file_path = f"/tmp/{server_name}.env" + env_content = "\n".join(f"{k}={v}" for k, v in env_vars.items()) + write_cmd = f"cat > {env_file_path} << 'ENVEOF'\n{env_content}\nENVEOF" + + result = session.container.exec_run( + ["sh", "-c", write_cmd], workdir=str(DockerComposeManager.CONTAINER_APP_DIR) + ) + if result.exit_code != 0: + logger.warning( + f"[MCP] Failed to write env file: {result.output.decode('utf-8')}" + ) + else: + log_harness_info( + logger, + task_name, + "agent", + f"Wrote env file with vars: {list(env_vars.keys())}", + ) + + # Build mcp add command + args_str = " ".join(mcp_config.args) + if env_file_path: + mcp_cmd = f"{agent_cli} mcp add {server_name} -- {mcp_config.command} --env-file {env_file_path} {args_str}" + else: + mcp_cmd = f"{agent_cli} mcp add {server_name} -- {mcp_config.command} {args_str}" + + result = session.container.exec_run( + ["sh", "-c", mcp_cmd], workdir=str(DockerComposeManager.CONTAINER_APP_DIR) + ) + + if result.exit_code != 0: + logger.warning( + f"[MCP] Server registration failed for {server_name}: " + f"{result.output.decode('utf-8')}" + ) + else: + log_harness_info( + logger, task_name, "agent", f"MCP server '{server_name}' configured" + ) + def perform_task( self, task_prompt: str, @@ -104,32 +202,9 @@ def perform_task( max_timeout_sec=config.setup_timeout_sec, # Use setup timeout for installation ) - # Optionally setup dbt MCP server - if self._use_mcp: - dbt_mcp_script = ( - Path(__file__).parent.parent.parent.parent - / "shared" - / "scripts" - / "setup-dbt-mcp.sh" - ) - session.copy_to_container( - dbt_mcp_script, - container_dir="/scripts", - container_filename="setup-dbt-mcp.sh", - ) - - # Pass db_type, project_type, and agent name - db_type = self._variant_config.get("db_type", "unknown") - project_type = self._variant_config.get("project_type", "unknown") - agent_name = self.NAME.value if hasattr(self.NAME, "value") else str(self.NAME) - session.send_keys( - [ - f"bash /scripts/setup-dbt-mcp.sh {db_type} {project_type} {agent_name}", - "Enter", - ], - block=True, - max_timeout_sec=config.setup_timeout_sec, - ) + # Configure MCP servers after agent is installed + if self._mcp_servers: + self._configure_mcp_servers(session, task_name) except TimeoutError: log_harness_info( logger, diff --git a/ade_bench/agents/installed_agents/claude_code/claude_code_agent.py b/ade_bench/agents/installed_agents/claude_code/claude_code_agent.py index d7746a87..72f9bfd9 100644 --- a/ade_bench/agents/installed_agents/claude_code/claude_code_agent.py +++ b/ade_bench/agents/installed_agents/claude_code/claude_code_agent.py @@ -15,7 +15,6 @@ class ClaudeCodeAgent(AbstractInstalledAgent): NAME = AgentName.CLAUDE_CODE - ALLOWED_TOOLS = ["Bash", "Edit", "Write", "NotebookEdit", "WebFetch", "mcp__dbt"] def __init__(self, **kwargs): super().__init__(**kwargs) @@ -40,7 +39,8 @@ def _run_agent_commands(self, task_prompt: str) -> list[TerminalCommand]: if self._model_name: command += f" --model {self._model_name}" - command += f" --allowedTools {' '.join(self.ALLOWED_TOOLS)}" + if self._allowed_tools: + command += f" --allowedTools {' '.join(self._allowed_tools)}" return [ TerminalCommand( @@ -62,10 +62,62 @@ def format_agent_log(self, log_path: Path) -> str | None: """ Format the Claude Code agent's log file into a human-readable string. + Also generates an HTML transcript at log_path.parent / "transcript.html" + using claude-code-transcripts if available. + + Args: log_path: Path to the raw agent.log file (JSON-lines format) Returns: Formatted log content as a string, or None if formatting failed """ + # Generate HTML transcript as a single well-known file + transcript_path = log_path.parent / "transcript.html" + self._log_formatter.generate_html_transcript(log_path, transcript_path) + + # Return text-formatted log return self._log_formatter.format_log(log_path) + + # Generic tools to filter out from tools_used reporting + _GENERIC_TOOLS = frozenset( + { + "Bash", + "Edit", + "Glob", + "Grep", + "Read", + "Write", + "WebFetch", + "WebSearch", + "Task", + "NotebookEdit", + "TodoRead", + "TodoWrite", + } + ) + + def extract_tools_used(self, log_path: Path) -> list[str] | None: + """ + Extract deduplicated tool names from Claude Code agent logs. + + Filters out generic tools (Bash, Edit, Glob, etc.) and expands + Skill tool calls to their actual skill names. + """ + try: + turns = self._log_formatter.parse_log_file(log_path) + tool_names = set() + for turn in turns: + for tool in turn.get("tools", []): + name = tool["name"] + # Expand Skill tool to actual skill name + if name == "Skill": + skill_name = tool.get("input", {}).get("skill") + if skill_name: + tool_names.add(f"skill:{skill_name}") + # Filter out generic tools + elif name not in self._GENERIC_TOOLS: + tool_names.add(name) + return sorted(tool_names) if tool_names else None + except Exception: + return None diff --git a/ade_bench/agents/installed_agents/claude_code/log_formatter.py b/ade_bench/agents/installed_agents/claude_code/log_formatter.py index e473a508..76e3837c 100644 --- a/ade_bench/agents/installed_agents/claude_code/log_formatter.py +++ b/ade_bench/agents/installed_agents/claude_code/log_formatter.py @@ -2,17 +2,23 @@ Log formatter for Claude Code agent. This module provides parsing and formatting utilities for Claude Code agent -log files (JSON-lines format). +log files (JSON-lines format), and generates HTML transcripts using +claude-code-transcripts. """ +import contextlib +import io import json +import logging import re -from io import StringIO +import tempfile from pathlib import Path from typing import Any, Dict, List from ade_bench.agents.log_formatter import LogFormatter +logger = logging.getLogger(__name__) + class ClaudeCodeLogFormatter(LogFormatter): """Log formatter for Claude Code agent JSON-lines format.""" @@ -23,6 +29,65 @@ def strip_ansi_codes(text: str) -> str: ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") return ansi_escape.sub("", text) + @staticmethod + def extract_jsonl_content(log_path: Path, inject_prompt: str | None = None) -> str: + """ + Extract only the JSON lines from a log file that may contain mixed content. + + The log file may contain terminal output before the JSON lines begin. + This method extracts only valid JSON lines. + + Claude Code's stream-json output doesn't include the initial user prompt, + only the assistant responses and tool results. If no user prompt with text + is found in the log, a synthetic one is injected so that transcript + generation tools can identify conversation boundaries. + + Args: + log_path: Path to the log file + inject_prompt: Optional prompt text to inject if none found + + Returns: + String containing only the JSON lines, newline-separated + """ + json_lines = [] + has_user_text_prompt = False + + with open(log_path, "r") as f: + for line in f: + stripped = line.strip() + if stripped.startswith("{"): + try: + # Validate it's actually JSON + data = json.loads(stripped) + json_lines.append(stripped) + + # Check if this is a user message with actual text content + if data.get("type") == "user": + content = data.get("message", {}).get("content", []) + if isinstance(content, str) and content.strip(): + has_user_text_prompt = True + elif isinstance(content, list): + for item in content: + if isinstance(item, dict) and item.get("type") == "text": + has_user_text_prompt = True + break + except json.JSONDecodeError: + continue + + # If no user prompt found, inject a synthetic one at the beginning + if not has_user_text_prompt and json_lines: + prompt_text = inject_prompt or "Claude Code Agent Session" + synthetic_prompt = json.dumps( + { + "type": "user", + "timestamp": "", + "message": {"role": "user", "content": prompt_text}, + } + ) + json_lines.insert(0, synthetic_prompt) + + return "\n".join(json_lines) + @staticmethod def format_tool_input(tool_name: str, tool_input: Dict[str, Any]) -> str: """Format tool input parameters nicely.""" @@ -165,7 +230,7 @@ def parse_log_file(self, log_path: Path) -> List[Dict[str, Any]]: def format_readable_log(self, turns: List[Dict[str, Any]]) -> str: """Format the parsed turns into a readable text string.""" - output = StringIO() + output = io.StringIO() output.write("=" * 80 + "\n") output.write("CLAUDE CODE AGENT INTERACTION LOG\n") @@ -221,3 +286,81 @@ def format_readable_log(self, turns: List[Dict[str, Any]]) -> str: output.write("=" * 80 + "\n") return output.getvalue() + + def generate_html_transcript(self, log_path: Path, output_path: Path) -> Path | None: + """ + Generate an HTML transcript using claude-code-transcripts. + + This method extracts JSON lines from the log file (which may contain + mixed terminal output and JSON) and uses claude-code-transcripts to + generate a clean HTML transcript at the specified output path. + + Args: + log_path: Path to the log file (may contain mixed content) + output_path: Desired output file path (e.g., sessions/transcript.html) + + Returns: + Path to the generated HTML file, or None if generation failed + """ + if not log_path.exists(): + logger.warning(f"Log file not found: {log_path}") + return None + + try: + import shutil + from claude_code_transcripts import generate_html + + # Extract only JSON lines from the log file + jsonl_content = self.extract_jsonl_content(log_path) + if not jsonl_content: + logger.warning(f"No JSON content found in {log_path}") + return None + + output_path.parent.mkdir(parents=True, exist_ok=True) + + # Use a temporary directory for claude-code-transcripts output, + # then copy the result to the single well-known output path + with tempfile.TemporaryDirectory() as tmp_output_dir: + tmp_output = Path(tmp_output_dir) + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".jsonl", delete=False + ) as tmp_file: + tmp_file.write(jsonl_content) + tmp_path = Path(tmp_file.name) + + try: + # Generate HTML transcript (suppress stdout/stderr from library) + with ( + contextlib.redirect_stdout(io.StringIO()), + contextlib.redirect_stderr(io.StringIO()), + ): + generate_html(tmp_path, tmp_output) + + # Find the generated file (index.html or page-001.html) + generated = None + for candidate in ["index.html", "page-001.html"]: + candidate_path = tmp_output / candidate + if candidate_path.exists(): + generated = candidate_path + break + + if generated is None: + logger.warning(f"No HTML output found in {tmp_output}") + return None + + # Copy to the well-known output path + shutil.copy2(generated, output_path) + return output_path + finally: + tmp_path.unlink(missing_ok=True) + + except ImportError: + logger.warning( + "claude-code-transcripts not installed. " + "Install with: pip install claude-code-transcripts" + ) + return None + except Exception as e: + logger.warning(f"Transcript generation failed: {e}") + return None diff --git a/ade_bench/agents/log_formatter.py b/ade_bench/agents/log_formatter.py index 7ce97140..7390d468 100644 --- a/ade_bench/agents/log_formatter.py +++ b/ade_bench/agents/log_formatter.py @@ -57,3 +57,19 @@ def format_log(self, log_path: Path) -> str | None: return self.format_readable_log(turns) except Exception: return None + + def generate_html_transcript(self, log_path: Path, output_path: Path) -> Path | None: + """ + Generate an HTML transcript from the log file. + + Override this method in subclasses to provide HTML transcript generation. + The default implementation returns None (no HTML transcript). + + Args: + log_path: Path to the log file to parse + output_path: Desired output file path (e.g., sessions/transcript.html) + + Returns: + Path to the generated HTML file, or None if not supported/failed + """ + return None diff --git a/ade_bench/cli/ab/main.py b/ade_bench/cli/ab/main.py index 80f82cd4..8ab2b27e 100644 --- a/ade_bench/cli/ab/main.py +++ b/ade_bench/cli/ab/main.py @@ -122,8 +122,10 @@ def run( help="Upload results to S3 bucket (bucket name is read from config)", ), log_level: str = typer.Option("INFO", "--log-level", help="Set the logging level"), - use_mcp: bool = typer.Option( - False, "--use-mcp", help="Enable MCP (Model Context Protocol) for the agent" + plugin_set: Optional[str] = typer.Option( + None, + "--plugin-set", + help="Plugin set names from plugin-sets.yaml, space-separated (default: use all default sets)", ), with_profiling: bool = typer.Option( False, @@ -210,7 +212,7 @@ def run( db_type=db, project_type=project_type, keep_alive=persist, - use_mcp=use_mcp, + plugin_set_names=plugin_set.split() if plugin_set else None, with_profiling=with_profiling, ) diff --git a/ade_bench/harness.py b/ade_bench/harness.py index 0b343742..ca0f3e41 100644 --- a/ade_bench/harness.py +++ b/ade_bench/harness.py @@ -25,9 +25,11 @@ from ade_bench.harness_models import ( BenchmarkResults, FailureMode, + PluginSet, RunMetadata, TrialResults, ) +from ade_bench.plugins.loader import PluginSetLoader from ade_bench.setup.setup_orchestrator import SetupOrchestrator from ade_bench.llms.base_llm import ContextLengthExceededError, ParseError from ade_bench.parsers.base_parser import UnitTestStatus, ParserResult @@ -65,7 +67,7 @@ def __init__( db_type: str | None = None, project_type: str | None = None, keep_alive: bool = False, - use_mcp: bool = False, + plugin_set_names: list[str] | None = None, with_profiling: bool = False, ): """ @@ -94,7 +96,7 @@ def __init__( db_type: Database type to filter variants (e.g., duckdb, postgres, sqlite, snowflake). project_type: Project type to filter variants (e.g., dbt, other). keep_alive: If True, keep containers alive when tasks fail for debugging. - use_mcp: If True, start a dbt MCP server after setup completes. + plugin_set_names: List of skill set names to use. If None, uses defaults. with_profiling: If True, will enable the cProfiler. """ self._run_uuid = None @@ -108,7 +110,9 @@ def __init__( self._db_filter = db_type self._project_type_filter = project_type self._keep_alive = keep_alive - self._use_mcp = use_mcp + self._plugin_set_names = plugin_set_names + self._plugin_sets: list[PluginSet] = [] + self._current_plugin_set: PluginSet | None = None self._with_profiling = with_profiling # Initialize setup orchestrator for variant-specific setup @@ -128,10 +132,9 @@ def __init__( self._s3_bucket = ade_bench_config.s3_bucket_name self._n_concurrent_trials = n_concurrent_trials self._n_attempts = n_attempts - self._run_path.mkdir(parents=True, exist_ok=True) self._init_dataset() - + self._init_plugin_sets() self._init_logger() @property @@ -181,8 +184,13 @@ def _create_agent_for_task(self, task_id: str) -> BaseAgent: if self._model_name: agent_kwargs["model_name"] = self._model_name - # Pass use_mcp flag to installed agents - agent_kwargs["use_mcp"] = self._use_mcp + # Pass allowed_tools from current plugin set + if self._current_plugin_set and self._current_plugin_set.allowed_tools: + agent_kwargs["allowed_tools"] = self._current_plugin_set.allowed_tools + + # Pass mcp_servers from current plugin set + if self._current_plugin_set and self._current_plugin_set.mcp_servers: + agent_kwargs["mcp_servers"] = self._current_plugin_set.mcp_servers return self._agent_factory.get_agent(**agent_kwargs) @@ -194,16 +202,44 @@ def _init_dataset(self) -> None: excluded_task_ids=self._exclude_task_ids, ) - def _init_logger(self) -> None: - file_handler = logging.FileHandler(self._log_output_path) - file_handler.setLevel(logging.DEBUG) - logger.addHandler(file_handler) + def _init_plugin_sets(self) -> None: + """Load and resolve plugin sets from configuration.""" + config_path = self._dataset_path.parent / "experiment_sets" / "plugin-sets.yaml" + if not config_path.exists(): + # No plugin sets config - use empty list (no plugins) + self._plugin_sets = [ + PluginSet( + name="no-plugins", + allowed_tools=["Bash", "Edit", "Write", "Read", "Glob", "Grep"], + ) + ] + return + loader = PluginSetLoader(config_path) + self._plugin_sets = loader.resolve_plugin_sets( + plugin_set_names=self._plugin_set_names, agent_name=self._agent_name.value + ) + + def _init_logger(self) -> None: + """Initialize console logging. File logging is initialized per plugin set.""" console_handler = logging.StreamHandler() console_handler.setLevel(self._log_level) logger.addHandler(console_handler) self._logger = logger.getChild(__name__) + self._file_handler: logging.FileHandler | None = None + + def _init_file_logger(self) -> None: + """Initialize or reinitialize file logging for the current run path.""" + # Remove existing file handler if present + if self._file_handler is not None: + logger.removeHandler(self._file_handler) + self._file_handler.close() + + # Create new file handler for the current run path + self._file_handler = logging.FileHandler(self._log_output_path) + self._file_handler.setLevel(logging.DEBUG) + logger.addHandler(self._file_handler) def _is_resolved(self, parser_result: ParserResult | None) -> bool: if parser_result is None: @@ -451,14 +487,21 @@ async def _run_agent_with_timeout( logging_dir: Path, agent: BaseAgent, task_name: str | None = None, + prompt_suffix: str = "", ) -> AgentResult | None: timeouts = TimeoutManager.get_timeouts_for_task(trial_handler.task) + + # Build the full prompt with optional suffix + full_prompt = trial_handler.task_prompt + if prompt_suffix: + full_prompt = f"{full_prompt}\n\n{prompt_suffix}" + loop = asyncio.get_event_loop() task = loop.run_in_executor( None, partial( agent.perform_task, - task_prompt=trial_handler.task_prompt, + task_prompt=full_prompt, session=session, logging_dir=logging_dir, task_name=task_name, @@ -476,6 +519,11 @@ def _run_agent( task_name: str | None = None, ) -> tuple[AgentResult | None, FailureMode]: try: + # Get prompt suffix from current plugin set + prompt_suffix = "" + if self._current_plugin_set and self._current_plugin_set.prompt_suffix: + prompt_suffix = self._current_plugin_set.prompt_suffix + result = asyncio.run( self._run_agent_with_timeout( trial_handler=trial_handler, @@ -483,6 +531,7 @@ def _run_agent( logging_dir=trial_handler.agent_logging_dir, agent=agent, task_name=task_name, + prompt_suffix=prompt_suffix, ) ) @@ -579,6 +628,7 @@ def _run_setup( session=session, file_diff_handler=file_diff_handler, trial_handler=trial_handler, + plugin_set=self._current_plugin_set, ) # Run setup with timeout using asyncio @@ -628,7 +678,18 @@ def _run_trial( model_name=self._model_name, db_type=config.get("db_type"), project_type=config.get("project_type"), - used_mcp=self._use_mcp, + plugin_set_name=self._current_plugin_set.name if self._current_plugin_set else None, + plugin_set_skills=( + self._current_plugin_set.skill_locations if self._current_plugin_set else None + ), + plugin_set_mcp_servers=( + list(self._current_plugin_set.mcp_servers.keys()) + if self._current_plugin_set + else None + ), + prompt_suffix=( + self._current_plugin_set.prompt_suffix if self._current_plugin_set else None + ), ) with spin_up_terminal( @@ -728,20 +789,21 @@ def _run_trial( parts = full_pane.split("=== ADE_BENCH_PHASE_DELIMITER_AGENT_START ===") post_agent_pane = parts[-1].strip() - # Try to generate a nicely formatted agent.txt from agent.log + # Write agent.log and attempt formatting via BaseAgent interface agent_log_path = trial_handler.sessions_path / "agent.log" formatted_content = None - - if agent_log_path.exists(): - try: - # Get formatted content from agent (returns string or None) - formatted_content = task_agent.format_agent_log(agent_log_path) - if formatted_content: - self._logger.debug( - "Generated formatted agent.txt from agent.log using agent's formatter" - ) - except Exception as e: - self._logger.warning(f"Failed to format agent.log: {e}. Using raw pane output.") + try: + agent_log_path.write_text(post_agent_pane) + # format_agent_log() returns None for agents without formatting (BaseAgent default) + formatted_content = task_agent.format_agent_log(agent_log_path) + if formatted_content: + self._logger.debug( + "Generated formatted agent.txt from agent.log using agent's formatter" + ) + except Exception as e: + self._logger.warning( + f"Failed to write/format agent.log: {e}. Using raw pane output." + ) # Write to file - either formatted content or fallback to raw pane if formatted_content: @@ -749,6 +811,13 @@ def _run_trial( else: trial_handler.agent_pane_path.write_text(post_agent_pane) + # Extract tools used (returns None for agents without tool extraction) + if agent_log_path.exists(): + try: + results.tools_used = task_agent.extract_tools_used(agent_log_path) + except Exception as e: + self._logger.debug(f"Could not extract tools used: {e}") + # Capture snapshot after agent and create agent diff if file_diff_handler: file_diff_handler.handle_phase_diffing( @@ -1269,7 +1338,18 @@ def _execute_single_trial( model_name=self._model_name, db_type=config.get("db_type"), project_type=config.get("project_type"), - used_mcp=self._use_mcp, + plugin_set_name=self._current_plugin_set.name if self._current_plugin_set else None, + plugin_set_skills=( + self._current_plugin_set.skill_locations if self._current_plugin_set else None + ), + plugin_set_mcp_servers=( + list(self._current_plugin_set.mcp_servers.keys()) + if self._current_plugin_set + else None + ), + prompt_suffix=( + self._current_plugin_set.prompt_suffix if self._current_plugin_set else None + ), ) return trial_results @@ -1377,10 +1457,35 @@ def _execute_tasks(self) -> BenchmarkResults: return results def run(self) -> BenchmarkResults: - """Run the harness. + """Run the benchmark with all configured plugin sets.""" + all_results = BenchmarkResults() + original_run_id = self._run_id + + for plugin_set in self._plugin_sets: + # Create run ID with plugin set suffix + self._run_id = f"{original_run_id}__{plugin_set.name}" + self._current_plugin_set = plugin_set + + # Ensure output directory exists for this plugin set + self._run_path.mkdir(parents=True, exist_ok=True) + + # Initialize file logger for this plugin set's run path + self._init_file_logger() + + self._logger.info(f"Starting run for plugin set: {plugin_set.name}") + + # Run trials for this plugin set + results = self._execute_trials() + all_results.results.extend(results.results) + + self._run_id = original_run_id + return all_results + + def _execute_trials(self) -> BenchmarkResults: + """Execute trials for the current skill set. Returns: - BenchmarkResults: The results of the harness run. + BenchmarkResults: The results of the trials. """ log_harness_info(logger, "system", "start", "STARTING HARNESS RUN") log_harness_info(logger, "system", "start", f"Run ID: {self._run_id}") diff --git a/ade_bench/harness_models.py b/ade_bench/harness_models.py index af3b0178..89a3ebb3 100644 --- a/ade_bench/harness_models.py +++ b/ade_bench/harness_models.py @@ -96,7 +96,12 @@ class TrialResults(BaseModel): model_name: str | None = None db_type: str | None = None project_type: str | None = None - used_mcp: bool | None = None + # Plugin set metadata + plugin_set_name: str | None = None + plugin_set_skills: list[str] | None = None + plugin_set_mcp_servers: list[str] | None = None + tools_used: list[str] | None = None + prompt_suffix: str | None = None class BenchmarkResults(BaseModel): @@ -229,3 +234,74 @@ def from_yaml_list(cls, path: Path) -> list["TerminalCommand"]: """Load a list of terminal commands from a YAML file.""" data = yaml.safe_load(path.read_text()) return [cls.model_validate(obj) for obj in data] + + +class McpServerConfig(BaseModel): + """Configuration for an MCP server.""" + + command: str + args: list[str] = [] + env: dict[str, str] = {} + + +class SkillOrigin(BaseModel): + """Configuration for a skill origin.""" + + location: str # Skill origin (e.g., git URL, local path, GitHub shorthand) + skill_names: list[str] = [] # Empty list means install all skills + + def install_all(self) -> bool: + """Return True if all skills should be installed from this origin.""" + return len(self.skill_names) == 0 + + +class PluginSet(BaseModel): + """Configuration for a set of plugins (skills and MCP servers).""" + + name: str + description: str = "" + default: bool = False + agents: list[str] | None = None # None = all agents compatible + skills: list[SkillOrigin] = [] + mcp_servers: dict[str, McpServerConfig] = {} + allowed_tools: list[str] = [] + prompt_suffix: str = "" # Appended to task prompt before execution + + def is_compatible_with_agent(self, agent_name: str) -> bool: + """Check if this plugin set is compatible with the given agent.""" + if self.agents is None: + return True + return agent_name in self.agents + + @property + def skill_locations(self) -> list[str]: + """Get list of skill locations as strings (for result tracking).""" + return [s.location for s in self.skills] + + +class PluginSetsConfig(BaseModel): + """Root configuration containing all plugin sets.""" + + sets: list[PluginSet] + + def get_defaults(self) -> list[PluginSet]: + """Get all plugin sets marked as default.""" + return [s for s in self.sets if s.default] + + def get_by_name(self, name: str) -> PluginSet | None: + """Get a plugin set by name.""" + for s in self.sets: + if s.name == name: + return s + return None + + def get_by_names(self, names: list[str]) -> list[PluginSet]: + """Get multiple plugin sets by name. Raises if any not found.""" + result = [] + for name in names: + plugin_set = self.get_by_name(name) + if plugin_set is None: + available = [s.name for s in self.sets] + raise ValueError(f"Unknown plugin set '{name}'. Available: {', '.join(available)}") + result.append(plugin_set) + return result diff --git a/ade_bench/plugins/__init__.py b/ade_bench/plugins/__init__.py new file mode 100644 index 00000000..a2a533f6 --- /dev/null +++ b/ade_bench/plugins/__init__.py @@ -0,0 +1,5 @@ +"""Plugin system for ADE-Bench.""" + +from .loader import PluginSetLoader + +__all__ = ["PluginSetLoader"] diff --git a/ade_bench/plugins/loader.py b/ade_bench/plugins/loader.py new file mode 100644 index 00000000..6c0c7a63 --- /dev/null +++ b/ade_bench/plugins/loader.py @@ -0,0 +1,75 @@ +"""Loader for plugin set configuration.""" + +from pathlib import Path +import yaml + +from ade_bench.harness_models import PluginSet, PluginSetsConfig + + +class PluginSetLoader: + """Loads and resolves plugin sets from YAML configuration.""" + + def __init__(self, config_path: Path): + self._config_path = config_path + self._config: PluginSetsConfig | None = None + + def load(self) -> PluginSetsConfig: + """Load the plugin sets configuration from YAML.""" + if not self._config_path.exists(): + raise FileNotFoundError(f"Plugin sets config not found: {self._config_path}") + + with open(self._config_path) as f: + data = yaml.safe_load(f) + + self._config = PluginSetsConfig(**data) + return self._config + + def resolve_plugin_sets( + self, + plugin_set_names: list[str] | None, + agent_name: str, + ) -> list[PluginSet]: + """Resolve which plugin sets to use for a run. + + Args: + plugin_set_names: Explicit plugin set names from --plugin-set, or None for defaults + agent_name: The agent being used (e.g., "claude", "gemini") + + Returns: + List of PluginSet objects to use + + Raises: + ValueError: If requested plugin set is not found or incompatible + """ + if self._config is None: + self.load() + + # Get plugin sets (explicit or defaults) + if plugin_set_names: + plugin_sets = self._config.get_by_names(plugin_set_names) + # Validate all are compatible with agent + for ps in plugin_sets: + if not ps.is_compatible_with_agent(agent_name): + raise ValueError( + f"Plugin set '{ps.name}' is not compatible with agent '{agent_name}'. " + f"Compatible agents: {ps.agents}" + ) + else: + plugin_sets = self._config.get_defaults() + + # Filter to compatible plugin sets + compatible = [ps for ps in plugin_sets if ps.is_compatible_with_agent(agent_name)] + + if not compatible: + if plugin_set_names: + raise ValueError( + f"No compatible plugin sets found for agent '{agent_name}' " + f"from requested: {plugin_set_names}" + ) + else: + raise ValueError( + f"No compatible plugin sets found for agent '{agent_name}'. " + f"No default plugin sets are compatible with this agent." + ) + + return compatible diff --git a/ade_bench/plugins/skills_handler.py b/ade_bench/plugins/skills_handler.py new file mode 100644 index 00000000..1b91dcfe --- /dev/null +++ b/ade_bench/plugins/skills_handler.py @@ -0,0 +1,62 @@ +"""Handler for installing skills via Vercel Skills CLI.""" + +import logging +from ade_bench.harness_models import PluginSet, SkillOrigin +from ade_bench.terminal.docker_compose_manager import DockerComposeManager + +logger = logging.getLogger(__name__) + + +class SkillsHandler: + """Installs skills from plugin set configuration.""" + + def install(self, plugin_set: PluginSet, terminal: DockerComposeManager) -> None: + """Install skills from the plugin set into the container. + + Args: + plugin_set: The plugin set configuration + terminal: The Docker container manager + """ + if not plugin_set.skills: + logger.debug(f"[SkillsHandler] No skills to install for '{plugin_set.name}'") + return + + for skill_origin in plugin_set.skills: + self._install_skill_origin(skill_origin, terminal) + + def _install_skill_origin( + self, skill_origin: SkillOrigin, terminal: DockerComposeManager + ) -> None: + """Install skills from a single skill origin. + + Args: + skill_origin: The skill origin configuration + terminal: The Docker container manager + """ + # Base command with non-interactive flags: + # -y: skip confirmation prompts + # -g: install globally (to ~/.agents/skills) + base_cmd = f"npx --yes skills add {skill_origin.location} -y -g" + + if skill_origin.install_all(): + cmd = f"{base_cmd} --all" + desc = f"all skills from {skill_origin.location}" + else: + # Use --skill flag for each skill name (per npx skills CLI syntax) + skill_flags = " ".join(f"--skill {name}" for name in skill_origin.skill_names) + cmd = f"{base_cmd} {skill_flags}" + desc = f"skills {skill_origin.skill_names} from {skill_origin.location}" + + logger.info(f"[SkillsHandler] Installing {desc}...") + + result = terminal.container.exec_run( + ["sh", "-c", cmd], workdir=str(DockerComposeManager.CONTAINER_APP_DIR) + ) + + if result.exit_code != 0: + logger.warning( + f"[SkillsHandler] Skills installation failed for {skill_origin.location}: " + f"{result.output.decode('utf-8')}" + ) + else: + logger.info(f"[SkillsHandler] Successfully installed {desc}") diff --git a/ade_bench/setup/agent_setup.py b/ade_bench/setup/agent_setup.py index eecd53aa..b9a34802 100644 --- a/ade_bench/setup/agent_setup.py +++ b/ade_bench/setup/agent_setup.py @@ -33,6 +33,7 @@ def setup_agent_config(terminal, task_id: str, trial_handler, logger) -> None: log_harness_info(logger, task_id, "setup", "Migrating agent config files...") + # Copy agent-specific config files if agent_name == AgentName.CLAUDE_CODE: _copy_config_file(terminal, trial_handler, "CLAUDE.md") elif agent_name == AgentName.GEMINI_CLI: diff --git a/ade_bench/setup/setup_orchestrator.py b/ade_bench/setup/setup_orchestrator.py index a6624579..4c886a40 100644 --- a/ade_bench/setup/setup_orchestrator.py +++ b/ade_bench/setup/setup_orchestrator.py @@ -1,5 +1,5 @@ """ -Simple setup orchestrator - just calls functions directly. +Setup orchestrator - coordinates task setup and plugin configuration. """ from typing import Dict, Any @@ -10,19 +10,29 @@ from .migration_setup import setup_migration from .agent_setup import setup_agent_config from ..utils.logger import log_harness_info +from ..harness_models import PluginSet +from ..plugins.skills_handler import SkillsHandler class SetupOrchestrator: - """Simple orchestrator that calls setup functions directly.""" + """Orchestrator that calls setup functions and configures plugins.""" def __init__( - self, logger=None, terminal=None, session=None, file_diff_handler=None, trial_handler=None + self, + logger=None, + terminal=None, + session=None, + file_diff_handler=None, + trial_handler=None, + plugin_set: PluginSet | None = None, ): self.logger = logger self.terminal = terminal self.session = session self.file_diff_handler = file_diff_handler self.trial_handler = trial_handler + self.plugin_set = plugin_set + self._skills_handler = SkillsHandler() def setup_task(self, task_id: str, variant: Dict[str, Any]) -> bool: """Setup a task for the given variant.""" @@ -43,6 +53,12 @@ def setup_task(self, task_id: str, variant: Dict[str, Any]) -> bool: # Logging is in the setup_agent_config function setup_agent_config(self.terminal, task_id, self.trial_handler, self.logger) + # Install skills if plugin set specified (MCP is configured after agent installation) + if self.plugin_set and self.plugin_set.skills: + log_harness_info(self.logger, task_id, "setup", "Installing skills...") + self._skills_handler.install(self.plugin_set, self.terminal) + log_harness_info(self.logger, task_id, "setup", "Skills installed") + # Set up the database db_type = variant.get("db_type") if db_type == "duckdb": diff --git a/ade_bench/utils/results_writer.py b/ade_bench/utils/results_writer.py index 54c68755..48834069 100644 --- a/ade_bench/utils/results_writer.py +++ b/ade_bench/utils/results_writer.py @@ -135,11 +135,13 @@ def write_results_tsv(results: BenchmarkResults, output_path: Path, run_id: str) "output_tokens", "cache_tokens", "turns", + "tools", "agent", "model_name", "db_type", "project_type", - "used_mcp", + "plugin_set", + "prompt_suffix", ] with open(output_path, "w", newline="") as f: @@ -164,6 +166,8 @@ def write_results_tsv(results: BenchmarkResults, output_path: Path, run_id: str) # Get failure type failure_type = get_failure_type(trial_result) + tools_str = ",".join(trial_result.tools_used) if trial_result.tools_used else "" + row = [ run_id, calc["task_id"], @@ -179,11 +183,13 @@ def write_results_tsv(results: BenchmarkResults, output_path: Path, run_id: str) calc["_output_tokens"], calc["_cache_tokens"], calc["_turns"], + tools_str, trial_result.agent or "", trial_result.model_name or "", trial_result.db_type or "", trial_result.project_type or "", - trial_result.used_mcp if trial_result.used_mcp is not None else "", + trial_result.plugin_set_name or "", + trial_result.prompt_suffix or "", ] writer.writerow(row) diff --git a/docker/base/Dockerfile.duckdb-dbt b/docker/base/Dockerfile.duckdb-dbt index 9dcab047..cd6b0a1c 100644 --- a/docker/base/Dockerfile.duckdb-dbt +++ b/docker/base/Dockerfile.duckdb-dbt @@ -1,6 +1,7 @@ FROM python:3.11-slim RUN apt-get update && apt-get install -y \ + git \ tmux asciinema \ curl \ && curl -fsSL https://deb.nodesource.com/setup_22.x | bash - \ diff --git a/docker/base/Dockerfile.snowflake-dbt b/docker/base/Dockerfile.snowflake-dbt index 7de7187b..2b72e832 100644 --- a/docker/base/Dockerfile.snowflake-dbt +++ b/docker/base/Dockerfile.snowflake-dbt @@ -1,6 +1,7 @@ FROM python:3.11-slim RUN apt-get update && apt-get install -y \ + git \ tmux asciinema \ curl \ && curl -sSL https://github.com/mikefarah/yq/releases/latest/download/yq_linux_amd64 \ diff --git a/docker/base/Dockerfile.snowflake-dbtf b/docker/base/Dockerfile.snowflake-dbtf index 4fd08b6b..d3d53269 100644 --- a/docker/base/Dockerfile.snowflake-dbtf +++ b/docker/base/Dockerfile.snowflake-dbtf @@ -1,6 +1,7 @@ FROM python:3.11-slim RUN apt-get update && apt-get install -y \ + git \ tmux asciinema \ curl \ && curl -sSL https://github.com/mikefarah/yq/releases/latest/download/yq_linux_amd64 \ diff --git a/experiment_sets/plugin-sets.yaml b/experiment_sets/plugin-sets.yaml new file mode 100644 index 00000000..fca75693 --- /dev/null +++ b/experiment_sets/plugin-sets.yaml @@ -0,0 +1,64 @@ +# Skill set configurations for ADE-Bench +# Use --plugin-set to select, or run without flag to use all defaults + +sets: + - name: all-dbt-skills + description: All dbt skills + default: false + skills: + - location: dbt-labs/dbt-agent-skills + mcp_servers: {} + allowed_tools: [Bash, Edit, Write, Read, Glob, Grep, Skill] + prompt_suffix: "You can use the dbt skills." + + - name: dbt-for-ae + description: Just the Using dbt for analytics engineering skill + default: false + skills: + - location: dbt-labs/dbt-agent-skills + skill_names: [using-dbt-for-analytics-engineering] + mcp_servers: {} + allowed_tools: [Bash, Edit, Write, Read, Glob, Grep, Skill] + prompt_suffix: "You can use the dbt skills." + + - name: dbt-mcp + description: dbt MCP server + default: false + skills: [] + mcp_servers: + dbt: + command: uvx + args: [dbt-mcp@latest] + env: + DISABLE_SEMANTIC_LAYER: "true" + DISABLE_DISCOVERY: "true" + DISABLE_ADMIN_API: "true" + DISABLE_SQL: "true" + DISABLE_DBT_CODEGEN: "true" + allowed_tools: [Bash, Edit, Write, Read, Glob, Grep, mcp__dbt__*] + prompt_suffix: "You can use the dbt MCP server." + + - name: dbt-skills-mcp + description: dbt skills and MCP server + default: false + skills: + - location: dbt-labs/dbt-agent-skills + mcp_servers: + dbt: + command: uvx + args: [dbt-mcp@latest] + env: + DISABLE_SEMANTIC_LAYER: "true" + DISABLE_DISCOVERY: "true" + DISABLE_ADMIN_API: "true" + DISABLE_SQL: "true" + DISABLE_DBT_CODEGEN: "true" + allowed_tools: [Bash, Edit, Write, Read, Glob, Grep, Skill, mcp__dbt__*] + prompt_suffix: "You can use the dbt skills and MCP server." + + - name: none + description: Baseline - no skills or MCP + default: true + skills: [] + mcp_servers: {} + allowed_tools: [Bash, Edit, Write, Read, Glob, Grep] diff --git a/pyproject.toml b/pyproject.toml index 6f5074e6..c33d546e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "duckdb>=0.9", "snowflake-connector-python>=3.0", "pyarrow>=10.0", + "claude-code-transcripts>=0.6", ] [project.optional-dependencies] diff --git a/scripts_python/analyze.py b/scripts_python/analyze.py index ae9a4a81..f74c3886 100755 --- a/scripts_python/analyze.py +++ b/scripts_python/analyze.py @@ -186,7 +186,7 @@ def get_canonical_column_order() -> List[str]: "model_name", "db_type", "project_type", - "used_mcp", + "plugin_set", ] diff --git a/scripts_python/create_sandbox.py b/scripts_python/create_sandbox.py index 650d2955..3db8c4ca 100644 --- a/scripts_python/create_sandbox.py +++ b/scripts_python/create_sandbox.py @@ -315,11 +315,6 @@ def main(): "--db", required=True, help="Database type (duckdb, sqlite, postgres, snowflake)" ) parser.add_argument("--project-type", required=True, help="Project type (dbt, other)") - parser.add_argument("--agent", required=False, help="Ignored") - parser.add_argument("--use-mcp", required=False, action="store_true", help="Ignored") - parser.add_argument("--persist", required=False, action="store_true", help="Ignored") - parser.add_argument("--no-diffs", required=False, action="store_true", help="Ignored") - parser.add_argument("--seed", required=False, action="store_true", help="Ignored") args = parser.parse_args() diff --git a/scripts_python/generate_results_html.py b/scripts_python/generate_results_html.py index 4e44b522..d32c0f3c 100644 --- a/scripts_python/generate_results_html.py +++ b/scripts_python/generate_results_html.py @@ -192,7 +192,7 @@ def _generate_summary_page(self, experiment_data: Dict[str, Any]): "{{ project_type }}", html.escape(summary["project_type"] or "Unknown") ) html_content = html_content.replace( - "{{ used_mcp }}", "Yes" if summary["used_mcp"] else "No" + "{{ plugin_set }}", html.escape(summary["plugin_set"] or "None") ) # Replace the entire table section with the generated HTML table @@ -332,34 +332,94 @@ def _generate_results_page( ) def _generate_panes_page(self, task_data: Dict[str, Any], task_dir: Path, task_html_dir: Path): - """Generate panes detail page.""" - panes_dir = task_dir / "panes" - content = "" - - if panes_dir.exists(): - pane_files = list(panes_dir.glob("*.txt")) - pane_dict = {p.name: p for p in pane_files} + """Generate panes detail page. - pane_order = ["pre-agent.txt", "agent.txt", "post-agent.txt"] + Displays terminal output in chronological order: + 1. pre-agent.txt (setup before agent runs) + 2. HTML transcript (if available) OR agent.txt (raw output) + 3. post-agent.txt (teardown after agent completes) + """ + import shutil - # Main panes in desired order, then any missing panes - main_panes = [pane_dict[p] for p in pane_order if p in pane_dict] - missing_panes = [p for p in pane_files if p.name not in pane_order] + panes_dir = task_dir / "panes" + transcript_path = task_dir / "sessions" / "transcript.html" - panes = main_panes + missing_panes + # Check for HTML transcript (single well-known file) + transcript_html = None + if transcript_path.exists(): + # Copy transcript file to HTML output + shutil.copy2(transcript_path, task_html_dir / "transcript.html") + transcript_html = "transcript.html" - for pane_file in panes: - content += f"\n{'='*80}\n" - content += f"FILE: {pane_file.name}\n" - content += f"{'='*80}\n\n" - with open(pane_file, "r") as f: - content += f.read() - content += "\n\n" - else: - content = "No panes directory found." + # Build content sections in chronological order + sections = [] - self._write_detail_page( - task_html_dir / "panes.html", "Terminal Panes", task_data["task_id"], content, "panes" + if panes_dir.exists(): + pane_dict = {p.name: p for p in panes_dir.glob("*.txt")} + + # 1. Pre-agent section + if "pre-agent.txt" in pane_dict: + with open(pane_dict["pre-agent.txt"], "r") as f: + pre_content = f.read().strip() + if pre_content: + sections.append(("Pre-Agent Setup", pre_content, "text")) + + # 2. Agent section - HTML transcript OR raw agent.txt + if transcript_html: + sections.append(("Agent Transcript", transcript_html, "iframe")) + elif "agent.txt" in pane_dict: + with open(pane_dict["agent.txt"], "r") as f: + agent_content = f.read().strip() + if agent_content: + sections.append(("Agent Output", agent_content, "text")) + + # 3. Post-agent section + if "post-agent.txt" in pane_dict: + with open(pane_dict["post-agent.txt"], "r") as f: + post_content = f.read().strip() + if post_content: + sections.append(("Post-Agent Output", post_content, "text")) + + # Any other pane files not in the standard order + other_panes = [ + p + for p in pane_dict.keys() + if p not in ["pre-agent.txt", "agent.txt", "post-agent.txt"] + ] + for pane_name in sorted(other_panes): + with open(pane_dict[pane_name], "r") as f: + other_content = f.read().strip() + if other_content: + sections.append((f"Other: {pane_name}", other_content, "text")) + + # Build HTML content from sections + content_parts = [] + for title, content, content_type in sections: + if content_type == "iframe": + content_parts.append( + f"""
+

{html.escape(title)}

+ +

Open transcript in new tab

+
""" + ) + else: + content_parts.append( + f"""
+

{html.escape(title)}

+ 
{html.escape(content)}
+
""" + ) + + if not content_parts: + content_parts.append("

No panes data found.

") + + # Write directly using template + self._write_panes_page( + task_html_dir / "panes.html", + "Terminal Panes", + task_data["task_id"], + "\n".join(content_parts), ) def _generate_diffs_page(self, task_data: Dict[str, Any], task_dir: Path, task_html_dir: Path): @@ -388,7 +448,12 @@ def _generate_diffs_page(self, task_data: Dict[str, Any], task_dir: Path, task_h ) def _write_detail_page( - self, output_path: Path, title: str, task_id: str, content: str, content_type: str + self, + output_path: Path, + title: str, + task_id: str, + content: str, + content_type: str, ): """Write a detail page using the template.""" template_path = self.templates_dir / "detail.html" @@ -408,6 +473,35 @@ def _write_detail_page( with open(output_path, "w") as f: f.write(html_content) + def _write_panes_page(self, output_path: Path, title: str, task_id: str, content_html: str): + """Write the panes page with pre-built HTML content. + + Unlike _write_detail_page which escapes content, this method accepts + raw HTML that has already been constructed with proper sections. + + Args: + output_path: Path to write the HTML file + title: Page title + task_id: Task identifier + content_html: Pre-built HTML content (not escaped) + """ + template_path = self.templates_dir / "detail.html" + if not template_path.exists(): + print(f"Error: Template not found: {template_path}") + return + + with open(template_path, "r") as f: + template_content = f.read() + + # Simple template replacement - content_html is already formatted + html_content = template_content.replace("{{ title }}", html.escape(title)) + html_content = html_content.replace("{{ task_id }}", html.escape(task_id)) + html_content = html_content.replace("{{ content }}", content_html) + html_content = html_content.replace("{{ content_type }}", "panes") + + with open(output_path, "w") as f: + f.write(html_content) + def main(): """Main function to generate HTML results.""" diff --git a/scripts_python/summarize_results.py b/scripts_python/summarize_results.py index cef3272e..25e83c26 100644 --- a/scripts_python/summarize_results.py +++ b/scripts_python/summarize_results.py @@ -1,3 +1,5 @@ +import html + from tabulate import tabulate from ade_bench.harness_models import BenchmarkResults from ade_bench.utils.results_writer import format_trial_result, get_failure_type, is_error_result @@ -97,6 +99,7 @@ def summarize_results(results: BenchmarkResults) -> Dict[str, Any]: "output_tokens": output_tokens_str, "cache_tokens": cache_tokens_str, "turns": turns_str, + "tools_used": result.tools_used or [], # Store numeric values for totals calculation "_tests_num": calc["_tests"], "_passed_num": calc["_tests_passed"], @@ -151,7 +154,7 @@ def summarize_results(results: BenchmarkResults) -> Dict[str, Any]: "inferred_model": inferred_model, "db_type": first_result.db_type if first_result else None, "project_type": first_result.project_type if first_result else None, - "used_mcp": first_result.used_mcp if first_result else None, + "plugin_set": first_result.plugin_set_name if first_result else None, "agent": first_result.agent if first_result else None, }, } @@ -211,7 +214,7 @@ def generate_html_table(results: BenchmarkResults) -> str: # Generate table with unique placeholders for action links and task button # Insert 'Task' as second column (after 'Task' id) - headers = [summary["headers"][0], "Task"] + summary["headers"][1:] + ["Actions"] + headers = [summary["headers"][0], "Task"] + summary["headers"][1:] + ["Tools", "Actions"] table_data = [] # Add task rows with unique placeholders @@ -230,6 +233,7 @@ def generate_html_table(results: BenchmarkResults) -> str: task["output_tokens"], task["cache_tokens"], task["turns"], + f"__TOOLS_{i}__", # Placeholder for tools list f"__ACTION_LINKS_{i}__", # Unique placeholder for action links ] table_data.append(row) @@ -250,6 +254,7 @@ def generate_html_table(results: BenchmarkResults) -> str: total_row["output_tokens"], total_row["cache_tokens"], total_row["turns"], + "", # No tools for total row "", # No action links for total row ] table_data.append(total_row_data) @@ -257,7 +262,7 @@ def generate_html_table(results: BenchmarkResults) -> str: # Generate the base table html_table = tabulate(table_data, headers=headers, tablefmt="html") - # Now replace the placeholders with actual action links and task buttons + # Now replace the placeholders with actual action links, task buttons, and tools for i, task in enumerate(summary["tasks"]): action_links = f'
Results Panes Diffs
' html_table = html_table.replace(f"__ACTION_LINKS_{i}__", action_links) @@ -265,6 +270,16 @@ def generate_html_table(results: BenchmarkResults) -> str: task_button = f'' html_table = html_table.replace(f"__TASK_BUTTON_{i}__", task_button) + # Format tools as comma-separated list with styled spans + tools_list = task.get("tools_used", []) + if tools_list: + tools_html = ", ".join( + f'{html.escape(tool)}' for tool in tools_list + ) + else: + tools_html = '-' + html_table = html_table.replace(f"__TOOLS_{i}__", tools_html) + return html_table diff --git a/shared/config/CLAUDE.md b/shared/config/CLAUDE.md index 7d41ec5d..aa2814ab 100644 --- a/shared/config/CLAUDE.md +++ b/shared/config/CLAUDE.md @@ -3,11 +3,13 @@ You are acting as an expert analyst and data engineer who is taksed with solving analytics and data engineering problems. Follow the requests given to you—do exactly what is asked, nothing more. ## Available Tools + - dbt: You have access to a dbt project, and its configuration files. The project may use dbt Fusion or standard dbt. - Snowflake: Each dbt project is connected to a Snowflake database. - dbt's MCP server: In some cases, you may have access to the dbt MCP server, which you should use when appropriate. ## Key Responsibilities + - Create a model → create the model file with appropriate config and SQL - Fix a bug → identify and fix the issue - Update a model → make the requested changes @@ -15,7 +17,8 @@ You are acting as an expert analyst and data engineer who is taksed with solving - Do not add extra work like creating tests, adding documentation, or refactoring code unless explicitly asked. ## Key Principles + - Do what's asked: Follow the specific request, no extras - Inspect data: Look at actual data to understand problems or find issues - Check your work: When necessary, validate your work by querying data or compiling the dbt models -- Use the MCP server: Consider using the MCP server to speed up or improve your work. \ No newline at end of file +- Use the MCP server: Consider using the MCP server to speed up or improve your work. diff --git a/shared/scripts/setup-dbt-mcp.sh b/shared/scripts/setup-dbt-mcp.sh deleted file mode 100755 index e564a3e1..00000000 --- a/shared/scripts/setup-dbt-mcp.sh +++ /dev/null @@ -1,77 +0,0 @@ -#!/bin/bash -# Setup dbt MCP server for installed agents -# Arguments: db_type project_type agent_name - -echo "Setting up dbt MCP server..." - -# Parse arguments -DB_TYPE="${1:-unknown}" -PROJECT_TYPE="${2:-unknown}" -AGENT_NAME="${3:-unknown}" - -echo "Database type: $DB_TYPE" -echo "Project type: $PROJECT_TYPE" -echo "Agent: $AGENT_NAME" - -# Check if project type is dbt -if [[ ! " dbt dbt-fusion " =~ " $PROJECT_TYPE " ]]; then - echo "Skipping dbt MCP setup - '$PROJECT_TYPE' is not supported" - exit 0 -fi - -# Check if database type is supported -if [[ ! " snowflake " =~ " $DB_TYPE " ]]; then - echo "Skipping dbt MCP setup - '$DB_TYPE' is not supported" - exit 0 -fi - -# Get working directory and env file location -project_dir=$(pwd) -env_file="${project_dir}/.env" - -# Find dbt path -dbt_path=$(which dbt) -if [ -z "$dbt_path" ]; then - echo "WARNING: dbt not found in PATH, skipping MCP setup" - exit 0 -fi - -# Create .env file for dbt-mcp -# TODO, because this probably a janky way to do this. -cat > "$env_file" << EOF -DBT_PROJECT_DIR=$project_dir -DBT_PATH=$dbt_path -DISABLE_DBT_CLI=false -DISABLE_SEMANTIC_LAYER=true -DISABLE_DISCOVERY=true -DISABLE_ADMIN_API=true -DISABLE_SQL=true -DISABLE_DBT_CODEGEN=true -EOF - -# Check if dbt-mcp is already installed (pre-installed in Docker image) -if ! command -v dbt-mcp &> /dev/null; then - echo "dbt-mcp not found, installing..." - uv tool install dbt-mcp --force - echo "dbt-mcp installed" -fi - -if [[ "$AGENT_NAME" == "claude" ]]; then - echo "Registering dbt MCP server with Claude..." - claude mcp add dbt -- uvx --env-file "$env_file" dbt-mcp - claude mcp list - -elif [[ "$AGENT_NAME" == "codex" ]]; then - echo "Registering dbt MCP server with Codex..." - codex mcp add dbt -- uvx --env-file "$env_file" dbt-mcp - codex mcp list - -elif [[ "$AGENT_NAME" == "gemini" ]]; then - echo "Registering dbt MCP server with Gemini..." - gemini mcp add dbt uvx -- --env-file "$env_file" dbt-mcp - gemini mcp list - -else - echo "Skipping dbt MCP setup - '$AGENT_NAME' is not supported" - exit 0 -fi \ No newline at end of file diff --git a/templates/summary.html b/templates/summary.html index f33c76a8..aa77a225 100644 --- a/templates/summary.html +++ b/templates/summary.html @@ -194,6 +194,22 @@ font-family: inherit; } + /* Tool tags styling */ + .tool-tag { + display: inline-block; + padding: 2px 6px; + background-color: #264f78; + color: #9cdcfe; + border-radius: 3px; + font-size: 10px; + margin: 1px; + white-space: nowrap; + } + + .no-tools { + color: #6a6a6a; + } + /* Modal/Dialog styles */ .modal-overlay { display: none; @@ -290,7 +306,7 @@

Run Configuration

Model: {{ model }}

DB Type: {{ db_type }}

Project Type: {{ project_type }}

-

MCP Enabled: {{ used_mcp }}

+

Plugin Set: {{ plugin_set }}

Results

diff --git a/tests/agents/__init__.py b/tests/agents/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/agents/installed_agents/__init__.py b/tests/agents/installed_agents/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/agents/installed_agents/test_abstract_installed_agent.py b/tests/agents/installed_agents/test_abstract_installed_agent.py new file mode 100644 index 00000000..aaff4aaf --- /dev/null +++ b/tests/agents/installed_agents/test_abstract_installed_agent.py @@ -0,0 +1,232 @@ +from unittest.mock import MagicMock +from pathlib import Path + +from ade_bench.agents.agent_name import AgentName +from ade_bench.agents.installed_agents.abstract_installed_agent import AbstractInstalledAgent +from ade_bench.harness_models import McpServerConfig, TerminalCommand + + +class ConcreteInstalledAgent(AbstractInstalledAgent): + """Concrete subclass for testing AbstractInstalledAgent.""" + + NAME = AgentName.CLAUDE_CODE + + @property + def _env(self) -> dict[str, str]: + return {"TEST_KEY": "test_value"} + + @property + def _install_agent_script(self) -> Path: + return Path("/fake/install.sh") + + def _run_agent_commands(self, task_prompt: str) -> list[TerminalCommand]: + return [] + + +def _make_session(exec_results=None): + """Create a mock TmuxSession with configurable exec_run results.""" + session = MagicMock() + if exec_results is None: + session.container.exec_run.return_value = MagicMock(exit_code=0, output=b"Success") + else: + session.container.exec_run.side_effect = exec_results + return session + + +class TestConfigureMcpServersNoop: + def test_no_mcp_servers(self): + """No calls when mcp_servers is empty.""" + agent = ConcreteInstalledAgent(mcp_servers={}) + session = _make_session() + agent._configure_mcp_servers(session, "test_task") + session.container.exec_run.assert_not_called() + + +class TestConfigureMcpServersBasic: + def test_single_server_with_env(self): + """Single MCP server writes env file and runs mcp add.""" + mcp_servers = { + "myserver": McpServerConfig( + command="uvx", + args=["some-mcp@latest"], + env={"FOO": "bar", "BAZ": "qux"}, + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + session = _make_session() + + agent._configure_mcp_servers(session, "test_task") + + # Should be called twice: once to write env file, once for mcp add + assert session.container.exec_run.call_count == 2 + + # First call: write env file + env_write_call = session.container.exec_run.call_args_list[0] + env_cmd = env_write_call[0][0] + assert env_cmd[0] == "sh" + assert env_cmd[1] == "-c" + assert "/tmp/myserver.env" in env_cmd[2] + assert "FOO=bar" in env_cmd[2] + assert "BAZ=qux" in env_cmd[2] + + # Second call: mcp add + mcp_add_call = session.container.exec_run.call_args_list[1] + mcp_cmd = mcp_add_call[0][0] + assert "mcp add myserver" in mcp_cmd[2] + assert "--env-file /tmp/myserver.env" in mcp_cmd[2] + assert "some-mcp@latest" in mcp_cmd[2] + + def test_server_without_env(self): + """MCP server with no env vars skips env file writing.""" + mcp_servers = { + "simple": McpServerConfig( + command="npx", + args=["simple-mcp"], + env={}, + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + session = _make_session() + + agent._configure_mcp_servers(session, "test_task") + + # Only 1 call (no env file to write, but dbt detection adds nothing) + assert session.container.exec_run.call_count == 1 + + mcp_cmd = session.container.exec_run.call_args[0][0] + assert "mcp add simple" in mcp_cmd[2] + assert "--env-file" not in mcp_cmd[2] + + +class TestConfigureMcpServersDbtDetection: + def test_dbt_server_by_name(self): + """dbt MCP server detected by server name 'dbt'.""" + mcp_servers = { + "dbt": McpServerConfig( + command="uvx", + args=["dbt-mcp@latest"], + env={"DISABLE_SQL": "true"}, + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + + # Mock exec_run: first for which dbt, second for env file, third for mcp add + which_dbt_result = MagicMock(exit_code=0, output=b"/usr/local/bin/dbt") + default_result = MagicMock(exit_code=0, output=b"Success") + session = _make_session( + exec_results=[ + which_dbt_result, # _get_dbt_dynamic_env: which dbt + default_result, # write env file + default_result, # mcp add + ] + ) + + agent._configure_mcp_servers(session, "test_task") + + # Verify which dbt was called + first_call = session.container.exec_run.call_args_list[0] + assert first_call[0][0] == ["sh", "-c", "which dbt"] + + # Verify env file includes dynamic vars + env_write_call = session.container.exec_run.call_args_list[1] + env_content = env_write_call[0][0][2] + assert "DISABLE_SQL=true" in env_content # static + assert "DBT_PROJECT_DIR=" in env_content # dynamic + assert "DBT_PATH=/usr/local/bin/dbt" in env_content # dynamic + assert "DISABLE_DBT_CLI=false" in env_content # dynamic + + def test_dbt_server_by_args(self): + """dbt MCP server detected by 'dbt-mcp' in args.""" + mcp_servers = { + "data-tools": McpServerConfig( + command="uvx", + args=["dbt-mcp@latest"], + env={}, + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + + which_dbt_result = MagicMock(exit_code=0, output=b"/usr/bin/dbt") + default_result = MagicMock(exit_code=0, output=b"Success") + session = _make_session( + exec_results=[ + which_dbt_result, # which dbt + default_result, # env file + default_result, # mcp add + ] + ) + + agent._configure_mcp_servers(session, "test_task") + + # Dynamic dbt env vars should be present since dbt-mcp is in args + env_write_call = session.container.exec_run.call_args_list[1] + env_content = env_write_call[0][0][2] + assert "DBT_PATH=/usr/bin/dbt" in env_content + + def test_static_env_takes_precedence(self): + """Static env vars from config are not overridden by dynamic vars.""" + mcp_servers = { + "dbt": McpServerConfig( + command="uvx", + args=["dbt-mcp@latest"], + env={"DISABLE_DBT_CLI": "true"}, # static override + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + + which_dbt_result = MagicMock(exit_code=0, output=b"/usr/bin/dbt") + default_result = MagicMock(exit_code=0, output=b"Success") + session = _make_session( + exec_results=[ + which_dbt_result, + default_result, + default_result, + ] + ) + + agent._configure_mcp_servers(session, "test_task") + + env_write_call = session.container.exec_run.call_args_list[1] + env_content = env_write_call[0][0][2] + # Static value should win over dynamic "false" + assert "DISABLE_DBT_CLI=true" in env_content + + +class TestConfigureMcpServersFailure: + def test_env_file_write_failure_logs_warning(self): + """Non-zero exit on env file write logs warning but continues.""" + mcp_servers = { + "myserver": McpServerConfig( + command="uvx", + args=["mcp@latest"], + env={"KEY": "val"}, + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + + env_fail = MagicMock(exit_code=1, output=b"Permission denied") + mcp_success = MagicMock(exit_code=0, output=b"OK") + session = _make_session(exec_results=[env_fail, mcp_success]) + + # Should not raise + agent._configure_mcp_servers(session, "test_task") + + # mcp add still called despite env file failure + assert session.container.exec_run.call_count == 2 + + def test_mcp_add_failure_logs_warning(self): + """Non-zero exit on mcp add logs warning but doesn't raise.""" + mcp_servers = { + "myserver": McpServerConfig( + command="uvx", + args=["mcp@latest"], + env={}, + ) + } + agent = ConcreteInstalledAgent(mcp_servers=mcp_servers) + + mcp_fail = MagicMock(exit_code=1, output=b"Command not found") + session = _make_session(exec_results=[mcp_fail]) + + # Should not raise + agent._configure_mcp_servers(session, "test_task") diff --git a/tests/plugins/__init__.py b/tests/plugins/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/plugins/test_loader.py b/tests/plugins/test_loader.py new file mode 100644 index 00000000..9ad30dfd --- /dev/null +++ b/tests/plugins/test_loader.py @@ -0,0 +1,144 @@ +import pytest +from pathlib import Path +from ade_bench.plugins.loader import PluginSetLoader +from ade_bench.harness_models import PluginSetsConfig + + +def test_loader_loads_yaml(tmp_path): + yaml_file = tmp_path / "plugin-sets.yaml" + yaml_file.write_text( + """ +sets: + - name: test + default: true + skills: [] + allowed_tools: [Bash] +""" + ) + loader = PluginSetLoader(yaml_file) + config = loader.load() + assert isinstance(config, PluginSetsConfig) + assert len(config.sets) == 1 + assert config.sets[0].name == "test" + + +def test_loader_file_not_found(): + loader = PluginSetLoader(Path("/nonexistent/plugin-sets.yaml")) + with pytest.raises(FileNotFoundError): + loader.load() + + +def test_loader_resolve_plugin_sets_explicit(): + """Explicit --plugin-set names are resolved.""" + yaml_content = """ +sets: + - name: a + default: false + allowed_tools: [Bash] + - name: b + default: true + allowed_tools: [Bash] +""" + import tempfile + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + f.flush() + loader = PluginSetLoader(Path(f.name)) + result = loader.resolve_plugin_sets(plugin_set_names=["a"], agent_name="claude") + assert len(result) == 1 + assert result[0].name == "a" + + +def test_loader_resolve_plugin_sets_defaults(): + """When no --plugin-set, use defaults.""" + yaml_content = """ +sets: + - name: a + default: false + allowed_tools: [Bash] + - name: b + default: true + allowed_tools: [Bash] + - name: c + default: true + allowed_tools: [Bash] +""" + import tempfile + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + f.flush() + loader = PluginSetLoader(Path(f.name)) + result = loader.resolve_plugin_sets(plugin_set_names=None, agent_name="claude") + assert len(result) == 2 + assert result[0].name == "b" + assert result[1].name == "c" + + +def test_loader_resolve_plugin_sets_filters_incompatible(): + """Plugin sets incompatible with agent are filtered out.""" + yaml_content = """ +sets: + - name: claude-only + default: true + agents: [claude] + allowed_tools: [Bash] + - name: all-agents + default: true + allowed_tools: [Bash] +""" + import tempfile + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + f.flush() + loader = PluginSetLoader(Path(f.name)) + + # Claude gets both + result = loader.resolve_plugin_sets(None, "claude") + assert len(result) == 2 + + # Gemini only gets all-agents + result = loader.resolve_plugin_sets(None, "gemini") + assert len(result) == 1 + assert result[0].name == "all-agents" + + +def test_loader_resolve_plugin_sets_error_on_incompatible_explicit(): + """Error when explicitly requested plugin set is incompatible.""" + yaml_content = """ +sets: + - name: claude-only + agents: [claude] + allowed_tools: [Bash] +""" + import tempfile + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + f.flush() + loader = PluginSetLoader(Path(f.name)) + + with pytest.raises(ValueError, match="not compatible with agent 'gemini'"): + loader.resolve_plugin_sets(["claude-only"], "gemini") + + +def test_loader_resolve_plugin_sets_error_when_none_compatible(): + """Error when no plugin sets are compatible with agent.""" + yaml_content = """ +sets: + - name: claude-only + default: true + agents: [claude] + allowed_tools: [Bash] +""" + import tempfile + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(yaml_content) + f.flush() + loader = PluginSetLoader(Path(f.name)) + + with pytest.raises(ValueError, match="No compatible plugin sets"): + loader.resolve_plugin_sets(None, "gemini") diff --git a/tests/plugins/test_skills_handler.py b/tests/plugins/test_skills_handler.py new file mode 100644 index 00000000..5e768131 --- /dev/null +++ b/tests/plugins/test_skills_handler.py @@ -0,0 +1,98 @@ +from unittest.mock import MagicMock +from ade_bench.plugins.skills_handler import SkillsHandler +from ade_bench.harness_models import PluginSet, SkillOrigin + + +def test_skills_handler_install_no_skills(): + """No-op when plugin set has no skills.""" + plugin_set = PluginSet(name="test", skills=[], allowed_tools=["Bash"]) + terminal = MagicMock() + + handler = SkillsHandler() + handler.install(plugin_set, terminal) + + terminal.container.exec_run.assert_not_called() + + +def test_skills_handler_install_all_skills(): + """Installs all skills from a repo when skill_names is empty.""" + plugin_set = PluginSet( + name="test", + skills=[SkillOrigin(location="dbt-labs/dbt-agent-skills")], + allowed_tools=["Bash"], + ) + terminal = MagicMock() + terminal.container.exec_run.return_value = MagicMock(exit_code=0, output=b"Success") + + handler = SkillsHandler() + handler.install(plugin_set, terminal) + + terminal.container.exec_run.assert_called_once() + call_args = terminal.container.exec_run.call_args + cmd = call_args[0][0] + assert "npx" in cmd[2] + assert "skills add" in cmd[2] + assert "dbt-labs/dbt-agent-skills" in cmd[2] + assert "--all" in cmd[2] + + +def test_skills_handler_install_specific_skills(): + """Installs only specified skills when skill_names is provided.""" + plugin_set = PluginSet( + name="test", + skills=[ + SkillOrigin( + location="dbt-labs/dbt-agent-skills", + skill_names=["using-dbt-for-analytics-engineering", "fetching-dbt-docs"], + ) + ], + allowed_tools=["Bash"], + ) + terminal = MagicMock() + terminal.container.exec_run.return_value = MagicMock(exit_code=0, output=b"Success") + + handler = SkillsHandler() + handler.install(plugin_set, terminal) + + terminal.container.exec_run.assert_called_once() + call_args = terminal.container.exec_run.call_args + cmd = call_args[0][0] + assert "npx" in cmd[2] + assert "skills add" in cmd[2] + assert "dbt-labs/dbt-agent-skills" in cmd[2] + # Verify --skill flag syntax is used for each skill + assert "--skill using-dbt-for-analytics-engineering" in cmd[2] + assert "--skill fetching-dbt-docs" in cmd[2] + assert "--all" not in cmd[2] + + +def test_skills_handler_install_multiple_origins(): + """Installs skills from multiple origins.""" + plugin_set = PluginSet( + name="test", + skills=[ + SkillOrigin(location="repo/a"), + SkillOrigin(location="repo/b"), + ], + allowed_tools=["Bash"], + ) + terminal = MagicMock() + terminal.container.exec_run.return_value = MagicMock(exit_code=0, output=b"Success") + + handler = SkillsHandler() + handler.install(plugin_set, terminal) + + assert terminal.container.exec_run.call_count == 2 + + +def test_skills_handler_install_failure_logs_warning(): + """Logs warning but doesn't raise on install failure.""" + plugin_set = PluginSet( + name="test", skills=[SkillOrigin(location="repo/failing")], allowed_tools=["Bash"] + ) + terminal = MagicMock() + terminal.container.exec_run.return_value = MagicMock(exit_code=1, output=b"npm ERR! not found") + + handler = SkillsHandler() + # Should not raise, just log warning + handler.install(plugin_set, terminal) diff --git a/tests/test_plugin_set.py b/tests/test_plugin_set.py new file mode 100644 index 00000000..579021cf --- /dev/null +++ b/tests/test_plugin_set.py @@ -0,0 +1,144 @@ +import pytest +from ade_bench.harness_models import PluginSet, McpServerConfig, PluginSetsConfig, SkillOrigin + + +def test_mcp_server_config_minimal(): + config = McpServerConfig(command="uvx", args=["dbt-mcp@latest"]) + assert config.command == "uvx" + assert config.args == ["dbt-mcp@latest"] + assert config.env == {} + + +def test_mcp_server_config_with_env(): + config = McpServerConfig(command="uvx", args=["dbt-mcp@latest"], env={"DISABLE_SQL": "true"}) + assert config.env == {"DISABLE_SQL": "true"} + + +def test_skill_origin_install_all(): + """Empty skill_names means install all skills.""" + origin = SkillOrigin(location="dbt-labs/dbt-agent-skills") + assert origin.install_all() is True + assert origin.skill_names == [] + + +def test_skill_origin_specific_skills(): + """Non-empty skill_names means install only those skills.""" + origin = SkillOrigin(location="dbt-labs/dbt-agent-skills", skill_names=["skill1", "skill2"]) + assert origin.install_all() is False + assert origin.skill_names == ["skill1", "skill2"] + + +def test_plugin_set_minimal(): + plugin_set = PluginSet(name="test", allowed_tools=["Bash"]) + assert plugin_set.name == "test" + assert plugin_set.description == "" + assert plugin_set.default is False + assert plugin_set.agents is None + assert plugin_set.skills == [] + assert plugin_set.mcp_servers == {} + assert plugin_set.allowed_tools == ["Bash"] + + +def test_plugin_set_full(): + plugin_set = PluginSet( + name="dbt-full", + description="Full dbt setup", + default=True, + agents=["claude"], + skills=[SkillOrigin(location="dbt-labs/dbt-agent-skills")], + mcp_servers={"dbt": McpServerConfig(command="uvx", args=["dbt-mcp@latest"])}, + allowed_tools=["Bash", "Skill", "mcp__dbt__*"], + ) + assert plugin_set.default is True + assert plugin_set.agents == ["claude"] + assert len(plugin_set.mcp_servers) == 1 + assert len(plugin_set.skills) == 1 + assert plugin_set.skills[0].location == "dbt-labs/dbt-agent-skills" + + +def test_plugin_set_is_compatible_with_agent_all(): + """When agents is None, compatible with all agents.""" + plugin_set = PluginSet(name="test", allowed_tools=["Bash"]) + assert plugin_set.is_compatible_with_agent("claude") is True + assert plugin_set.is_compatible_with_agent("gemini") is True + + +def test_plugin_set_is_compatible_with_agent_restricted(): + """When agents is set, only compatible with listed agents.""" + plugin_set = PluginSet(name="test", agents=["claude"], allowed_tools=["Bash"]) + assert plugin_set.is_compatible_with_agent("claude") is True + assert plugin_set.is_compatible_with_agent("gemini") is False + + +def test_plugin_sets_config_from_yaml(): + yaml_content = """ +sets: + - name: no-plugins + default: true + skills: [] + allowed_tools: [Bash, Read] + - name: dbt-skills + agents: [claude] + skills: + - location: dbt-labs/dbt-agent-skills + skill_names: [] + allowed_tools: [Bash, Skill] +""" + import yaml + + data = yaml.safe_load(yaml_content) + config = PluginSetsConfig(**data) + assert len(config.sets) == 2 + assert config.sets[0].name == "no-plugins" + assert config.sets[0].default is True + assert config.sets[1].skills[0].location == "dbt-labs/dbt-agent-skills" + + +def test_plugin_sets_config_get_defaults(): + config = PluginSetsConfig( + sets=[ + PluginSet(name="a", default=True, allowed_tools=["Bash"]), + PluginSet(name="b", default=False, allowed_tools=["Bash"]), + PluginSet(name="c", default=True, allowed_tools=["Bash"]), + ] + ) + defaults = config.get_defaults() + assert len(defaults) == 2 + assert defaults[0].name == "a" + assert defaults[1].name == "c" + + +def test_plugin_sets_config_get_by_name(): + config = PluginSetsConfig( + sets=[ + PluginSet(name="a", allowed_tools=["Bash"]), + PluginSet(name="b", allowed_tools=["Bash"]), + ] + ) + assert config.get_by_name("a").name == "a" + assert config.get_by_name("b").name == "b" + assert config.get_by_name("nonexistent") is None + + +def test_plugin_sets_config_get_by_names(): + config = PluginSetsConfig( + sets=[ + PluginSet(name="a", allowed_tools=["Bash"]), + PluginSet(name="b", allowed_tools=["Bash"]), + PluginSet(name="c", allowed_tools=["Bash"]), + ] + ) + result = config.get_by_names(["a", "c"]) + assert len(result) == 2 + assert result[0].name == "a" + assert result[1].name == "c" + + +def test_plugin_sets_config_get_by_names_unknown_raises(): + config = PluginSetsConfig( + sets=[ + PluginSet(name="a", allowed_tools=["Bash"]), + ] + ) + with pytest.raises(ValueError, match="Unknown plugin set"): + config.get_by_names(["a", "nonexistent"]) diff --git a/uv.lock b/uv.lock index d91553a2..0cdfba50 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.12'", @@ -14,6 +14,7 @@ source = { editable = "." } dependencies = [ { name = "anthropic" }, { name = "boto3" }, + { name = "claude-code-transcripts" }, { name = "click" }, { name = "docker" }, { name = "duckdb" }, @@ -50,6 +51,7 @@ requires-dist = [ { name = "anthropic", specifier = ">=0.18" }, { name = "black", marker = "extra == 'dev'", specifier = ">=23.0" }, { name = "boto3", specifier = ">=1.26" }, + { name = "claude-code-transcripts", specifier = ">=0.6" }, { name = "click", specifier = ">=8.1" }, { name = "docker", specifier = ">=6.0" }, { name = "duckdb", specifier = ">=0.9" }, @@ -317,6 +319,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/20/94/c5790835a017658cbfabd07f3bfb549140c3ac458cfc196323996b10095a/charset_normalizer-3.4.2-py3-none-any.whl", hash = "sha256:7f56930ab0abd1c45cd15be65cc741c28b1c9a34876ce8c17a2fa107810c0af0", size = 52626, upload-time = "2025-05-02T08:34:40.053Z" }, ] +[[package]] +name = "claude-code-transcripts" +version = "0.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "click-default-group" }, + { name = "httpx" }, + { name = "jinja2" }, + { name = "markdown" }, + { name = "questionary" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/5e/28056099926dcf7432d9274865fca074736522ff6ad1f1a59312dcd0bb69/claude_code_transcripts-0.6.tar.gz", hash = "sha256:c4cf35cd7f02bf6b71cb577c0a39ad25a184eeaf13ec6a2c8d6cdc2900324390", size = 29011, upload-time = "2026-01-25T05:50:39.443Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/b9/4422e5c8d40988810721a6fa6e48139b242652319c458f1f76408ef3cf54/claude_code_transcripts-0.6-py3-none-any.whl", hash = "sha256:b0d5afbd0355ff7eefcfd64cb9f248d443d6d9899d0997d6fec4256b4dcb9a19", size = 32576, upload-time = "2026-01-25T05:50:38.201Z" }, +] + [[package]] name = "click" version = "8.2.1" @@ -329,6 +348,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/85/32/10bb5764d90a8eee674e9dc6f4db6a0ab47c8c4d0d83c27f7c39ac415a4d/click-8.2.1-py3-none-any.whl", hash = "sha256:61a3265b914e850b85317d0b3109c7f8cd35a670f963866005d6ef1d5175a12b", size = 102215, upload-time = "2025-05-20T23:19:47.796Z" }, ] +[[package]] +name = "click-default-group" +version = "1.2.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/ce/edb087fb53de63dad3b36408ca30368f438738098e668b78c87f93cd41df/click_default_group-1.2.4.tar.gz", hash = "sha256:eb3f3c99ec0d456ca6cd2a7f08f7d4e91771bef51b01bdd9580cc6450fe1251e", size = 3505, upload-time = "2023-08-04T07:54:58.425Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/1a/aff8bb287a4b1400f69e09a53bd65de96aa5cee5691925b38731c67fc695/click_default_group-1.2.4-py2.py3-none-any.whl", hash = "sha256:9b60486923720e7fc61731bdb32b617039aba820e22e1c88766b1125592eaa5f", size = 4123, upload-time = "2023-08-04T07:54:56.875Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -622,6 +653,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d6/2f/9d207039fcfa00d3b30e4d765f062fbcc42c873c7518a8cfebb3eafd00e0/libtmux-0.46.2-py3-none-any.whl", hash = "sha256:6c32dbf22bde8e5e33b2714a4295f6e838dc640f337cd4c085a044f6828c7793", size = 60873, upload-time = "2025-05-26T19:40:02.284Z" }, ] +[[package]] +name = "markdown" +version = "3.10.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b7/b1/af95bcae8549f1f3fd70faacb29075826a0d689a27f232e8cee315efa053/markdown-3.10.1.tar.gz", hash = "sha256:1c19c10bd5c14ac948c53d0d762a04e2fa35a6d58a6b7b1e6bfcbe6fefc0001a", size = 365402, upload-time = "2026-01-21T18:09:28.206Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/1b/6ef961f543593969d25b2afe57a3564200280528caa9bd1082eecdd7b3bc/markdown-3.10.1-py3-none-any.whl", hash = "sha256:867d788939fe33e4b736426f5b9f651ad0c0ae0ecf89df0ca5d1176c70812fe3", size = 107684, upload-time = "2026-01-21T18:09:27.203Z" }, +] + [[package]] name = "markdown-it-py" version = "3.0.0" @@ -928,6 +968,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0a/9e/8c62f05b104d9f00edbb4c298b152deceb393ea67f0288d89d1139d7a859/podman-5.6.0-py3-none-any.whl", hash = "sha256:967ff8ad8c6b851bc5da1a9410973882d80e235a9410b7d1e931ce0c3324fbe3", size = 88713, upload-time = "2025-09-05T09:42:38.405Z" }, ] +[[package]] +name = "prompt-toolkit" +version = "3.0.52" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/96/06e01a7b38dce6fe1db213e061a4602dd6032a8a97ef6c1a862537732421/prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855", size = 434198, upload-time = "2025-08-27T15:24:02.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, +] + [[package]] name = "psycopg" version = "3.2.9" @@ -1248,6 +1300,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/de/02b54f42487e3d3c6efb3f89428677074ca7bf43aae402517bc7cca949f3/PyYAML-6.0.2-cp313-cp313-win_amd64.whl", hash = "sha256:8388ee1976c416731879ac16da0aff3f63b286ffdd57cdeb95f3f2e085687563", size = 156446, upload-time = "2024-08-06T20:33:04.33Z" }, ] +[[package]] +name = "questionary" +version = "2.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "prompt-toolkit" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f6/45/eafb0bba0f9988f6a2520f9ca2df2c82ddfa8d67c95d6625452e97b204a5/questionary-2.1.1.tar.gz", hash = "sha256:3d7e980292bb0107abaa79c68dd3eee3c561b83a0f89ae482860b181c8bd412d", size = 25845, upload-time = "2025-08-28T19:00:20.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/26/1062c7ec1b053db9e499b4d2d5bc231743201b74051c973dadeac80a8f43/questionary-2.1.1-py3-none-any.whl", hash = "sha256:a51af13f345f1cdea62347589fbb6df3b290306ab8930713bfae4d475a7d4a59", size = 36753, upload-time = "2025-08-28T19:00:19.56Z" }, +] + [[package]] name = "requests" version = "2.32.3" @@ -1585,3 +1649,12 @@ sdist = { url = "https://files.pythonhosted.org/packages/8a/78/16493d9c386d8e60e wheels = [ { url = "https://files.pythonhosted.org/packages/6b/11/cc635220681e93a0183390e26485430ca2c7b5f9d33b15c74c2861cb8091/urllib3-2.4.0-py3-none-any.whl", hash = "sha256:4e16665048960a0900c702d4a66415956a584919c03361cac9f1df5c5dd7e813", size = 128680, upload-time = "2025-04-10T15:23:37.377Z" }, ] + +[[package]] +name = "wcwidth" +version = "0.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c2/62/a7c072fbfefb2980a00f99ca994279cb9ecf310cb2e6b2a4d2a28fe192b3/wcwidth-0.5.3.tar.gz", hash = "sha256:53123b7af053c74e9fe2e92ac810301f6139e64379031f7124574212fb3b4091", size = 157587, upload-time = "2026-01-31T03:52:10.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/c1/d73f12f8cdb1891334a2ccf7389eed244d3941e74d80dd220badb937f3fb/wcwidth-0.5.3-py3-none-any.whl", hash = "sha256:d584eff31cd4753e1e5ff6c12e1edfdb324c995713f75d26c29807bb84bf649e", size = 92981, upload-time = "2026-01-31T03:52:09.14Z" }, +]