wafer-cli 0.2.23__py3-none-any.whl → 0.2.25__py3-none-any.whl

wafer/tests/test_eval_cli_parity.py

@@ -0,0 +1,199 @@
+"""Test that eval and CLI agent configs are in sync.
+
+Run: python -m pytest apps/wafer-cli/wafer/tests/test_eval_cli_parity.py -v
+  or: python apps/wafer-cli/wafer/tests/test_eval_cli_parity.py  (standalone)
+  or: python apps/wafer-cli/wafer/tests/test_eval_cli_parity.py --dump  (side-by-side)
+
+Checks:
+- Same bash allowlist (identity check — must be the same object)
+- Same enabled tools
+- Same system prompt text (modulo runtime variables like pool name)
+- CLI instructions generated from same allowlist
+
+Coverage notes:
+  These tests verify config-level parity. After composition, both paths
+  flow through the same codepath (Actor → run_agent_step → rollout →
+  anthropic.py provider) with no further system prompt modifications.
+
+  Known infra-level differences NOT tested here:
+  - Claude Code identity prefix: anthropic.py prepends "You are Claude Code..."
+    when using OAuth/Claude Code API keys. Eval typically uses raw API keys,
+    so eval agents may not get this prefix. This is an auth concern, not a
+    prompt content concern.
+  - Skills layer: CLI path (wevin_cli.py) can append skill metadata if
+    template.include_skills is True. Currently False for optimize-kernelbench,
+    so this is a no-op. Eval path doesn't have this layer at all.
+"""
+
+from __future__ import annotations
+
+import difflib
+import string
+import sys
+
+
+def _normalize_prompt(prompt: str) -> str:
+    """Normalize runtime variables so eval and CLI prompts are comparable."""
+    # Replace pool/target specifics with placeholders
+    prompt = prompt.replace("--pool mi300x-pool", "--pool POOL")
+    prompt = prompt.replace("--pool kernelbench-pool", "--pool POOL")
+    prompt = prompt.replace("--target mi300x", "--target TARGET")
+    return prompt
+
+
+def test_bash_allowlist_parity() -> None:
+    from wafer.agent_defaults import KERNELBENCH_BASH_ALLOWLIST
+    from wafer.templates.optimize_kernelbench import template
+
+    # Template should use the shared allowlist (same object)
+    assert template.bash_allowlist is KERNELBENCH_BASH_ALLOWLIST, (
+        "Template bash_allowlist is not the shared KERNELBENCH_BASH_ALLOWLIST. "
+        "Import it from wafer.agent_defaults instead of defining a local copy."
+    )
+
+    # Eval should also alias to the shared allowlist
+    from optimize_kernelbench_eval.base_config import BASH_ALLOWLIST
+
+    assert BASH_ALLOWLIST is KERNELBENCH_BASH_ALLOWLIST, (
+        "Eval BASH_ALLOWLIST is not the shared KERNELBENCH_BASH_ALLOWLIST. "
+        "Import it from wafer.agent_defaults instead of defining a local copy."
+    )
+
+
+def test_enabled_tools_parity() -> None:
+    from wafer.agent_defaults import ENABLED_TOOLS
+    from wafer.templates.optimize_kernelbench import template
+
+    assert template.tools == ENABLED_TOOLS, (
+        f"Template tools {template.tools} != shared ENABLED_TOOLS {ENABLED_TOOLS}"
+    )
+
+
+def test_system_prompt_parity() -> None:
+    """The task-specific system prompt should be identical between eval and CLI
+    (after normalizing runtime variables like pool name)."""
+    from optimize_kernelbench_eval.base_config import SYSTEM_PROMPT as EVAL_PROMPT
+
+    from wafer.templates.optimize_kernelbench import template
+
+    # Format eval prompt with HIP defaults (most common)
+    eval_formatted = EVAL_PROMPT.format(
+        backend="HIP",
+        backend_lower="hip",
+        target_flag="--pool POOL",
+        reference_path="<reference_file>",
+    )
+
+    # Format CLI template prompt with matching defaults
+    params = dict(template.defaults)
+    params["target_flag"] = "--pool POOL"
+    cli_formatted = string.Template(template.system_prompt).safe_substitute(**params)
+
+    eval_normalized = _normalize_prompt(eval_formatted)
+    cli_normalized = _normalize_prompt(cli_formatted)
+
+    if eval_normalized != cli_normalized:
+        diff = "\n".join(
+            difflib.unified_diff(
+                eval_normalized.splitlines(),
+                cli_normalized.splitlines(),
+                fromfile="eval (base_config.py)",
+                tofile="cli (optimize_kernelbench.py template)",
+                lineterm="",
+                n=2,
+            )
+        )
+        raise AssertionError(
+            f"System prompts differ between eval and CLI template:\n\n{diff}\n\n"
+            "Both should define the same task instructions. "
+            "Edit one to match the other."
+        )
+
+
+def test_cli_instructions_identical() -> None:
+    """Both paths should generate the same CLI instructions
+    (since they use the same bash_allowlist)."""
+    from wafer.agent_defaults import KERNELBENCH_BASH_ALLOWLIST
+    from wafer.cli_instructions import build_cli_instructions
+    from wafer.templates.optimize_kernelbench import template
+
+    eval_instructions = build_cli_instructions(KERNELBENCH_BASH_ALLOWLIST)
+    cli_instructions = build_cli_instructions(template.bash_allowlist)
+
+    assert eval_instructions == cli_instructions, (
+        "CLI instructions differ — this means the bash allowlists diverged."
+    )
+    assert len(eval_instructions) > 0, "CLI instructions should not be empty"
+
+
+def _dump_full_prompts() -> None:
+    """Standalone: dump both composed prompts for manual comparison."""
+    from optimize_kernelbench_eval.base_config import SYSTEM_PROMPT as EVAL_PROMPT
+
+    from wafer.agent_defaults import KERNELBENCH_BASH_ALLOWLIST
+    from wafer.cli_instructions import build_cli_instructions
+    from wafer.templates.optimize_kernelbench import template
+
+    cli_instructions = build_cli_instructions(KERNELBENCH_BASH_ALLOWLIST)
+
+    # Eval
+    eval_sys = EVAL_PROMPT.format(
+        backend="HIP",
+        backend_lower="hip",
+        target_flag="--pool mi300x-pool",
+        reference_path="<reference_file>",
+    )
+    eval_sys += "\n\n" + cli_instructions
+
+    # CLI
+    params = dict(template.defaults)
+    cli_sys = string.Template(template.system_prompt).safe_substitute(**params)
+    cli_sys += "\n\n" + build_cli_instructions(template.bash_allowlist)
+
+    print("=" * 60)
+    print("EVAL SYSTEM PROMPT")
+    print("=" * 60)
+    print(eval_sys)
+    print()
+    print("=" * 60)
+    print("CLI SYSTEM PROMPT")
+    print("=" * 60)
+    print(cli_sys)
+    print()
+
+    # Diff
+    eval_norm = _normalize_prompt(eval_sys)
+    cli_norm = _normalize_prompt(cli_sys)
+    diff = list(
+        difflib.unified_diff(
+            eval_norm.splitlines(),
+            cli_norm.splitlines(),
+            fromfile="eval",
+            tofile="cli",
+            lineterm="",
+            n=1,
+        )
+    )
+    if diff:
+        print("=" * 60)
+        print("DIFFERENCES (after normalizing pool names)")
+        print("=" * 60)
+        for line in diff:
+            print(line)
+    else:
+        print("IDENTICAL (after normalizing pool names)")
+
+
+if __name__ == "__main__":
+    if "--dump" in sys.argv:
+        _dump_full_prompts()
+    else:
+        test_bash_allowlist_parity()
+        print("PASS: bash_allowlist_parity")
+        test_enabled_tools_parity()
+        print("PASS: enabled_tools_parity")
+        test_system_prompt_parity()
+        print("PASS: system_prompt_parity")
+        test_cli_instructions_identical()
+        print("PASS: cli_instructions_identical")
+        print("\nAll parity checks passed.")

wafer/trace_compare.py

@@ -0,0 +1,183 @@
+"""CLI wrapper for trace comparison commands.
+
+This module provides the CLI interface for the `wafer compare` commands.
+All core logic is in wafer_core.lib.trace_compare.
+"""
+
+import sys
+from pathlib import Path
+
+import typer
+
+from wafer_core.lib.trace_compare import (
+    analyze_fusion_differences,
+    analyze_traces,
+    format_csv,
+    format_fusion_csv,
+    format_fusion_json,
+    format_fusion_text,
+    format_json,
+    format_text,
+)
+
+
+def compare_traces(
+    trace1: Path,
+    trace2: Path,
+    output: Path | None = None,
+    output_format: str = "text",
+    phase: str = "all",
+    show_layers: bool = False,
+    show_all: bool = False,
+    show_stack_traces: bool = False,
+) -> None:
+    """Compare two GPU traces and generate performance report.
+
+    Args:
+        trace1: Path to first trace file (AMD or NVIDIA)
+        trace2: Path to second trace file (AMD or NVIDIA)
+        output: Optional output file path (default: stdout)
+        output_format: Output format ('text', 'text-layers', 'csv', 'csv-layers', or 'json')
+        phase: Filter by phase ('all', 'prefill', or 'decode')
+        show_layers: Show layer-wise performance breakdown (text format only)
+        show_all: Show all items without truncation (applies to layers, operations, kernels)
+        show_stack_traces: Show Python stack traces for operations
+    """
+    # Validate files exist
+    if not trace1.exists():
+        typer.secho(f"❌ File not found: {trace1}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    if not trace2.exists():
+        typer.secho(f"❌ File not found: {trace2}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Analyze traces
+    # Only show progress messages for non-JSON formats (JSON needs clean stdout)
+    if output_format != 'json':
+        typer.echo("📊 Loading traces...")
+
+    # Determine how many stack traces to collect
+    max_stacks = 0 if (show_stack_traces and show_all) else (3 if show_stack_traces else 3)
+
+    try:
+        results = analyze_traces(
+            trace1,
+            trace2,
+            phase_filter=phase,
+            max_stacks=max_stacks,
+        )
+    except ValueError as e:
+        typer.secho(f"❌ {e}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+    except Exception as e:
+        typer.secho(f"❌ Error analyzing traces: {e}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Show loading confirmation
+    if output_format != 'json':
+        meta = results["metadata"]
+        # Determine which trace is AMD and which is NVIDIA
+        if meta['trace1_platform'] == 'AMD':
+            amd_gpu, nvidia_gpu = meta['trace1_gpu'], meta['trace2_gpu']
+        else:
+            amd_gpu, nvidia_gpu = meta['trace2_gpu'], meta['trace1_gpu']
+        typer.echo(f"✅ Loaded: AMD ({amd_gpu}) vs NVIDIA ({nvidia_gpu})")
+    typer.echo()
+
+    # Generate output based on format
+    if output_format == "text":
+        output_str = format_text(results, show_layers=show_layers, show_all=show_all, show_stack_traces=show_stack_traces)
+    elif output_format == "text-layers":
+        output_str = format_text(results, show_layers=True, show_all=show_all, show_stack_traces=show_stack_traces)
+    elif output_format == "csv":
+        output_str = format_csv(results, report_type="operations")
+    elif output_format == "csv-layers":
+        output_str = format_csv(results, report_type="layers")
+    elif output_format == "json":
+        output_str = format_json(results)
+    else:
+        typer.secho(f"❌ Unknown format: {output_format}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Write output
+    if output:
+        output.write_text(output_str)
+        typer.secho(f"✅ Report saved to {output}", fg=typer.colors.GREEN)
+    else:
+        typer.echo(output_str)
+
+
+def compare_fusion(
+    trace1: Path,
+    trace2: Path,
+    output: Path | None = None,
+    format_type: str = "text",
+    min_group_size: int = 50,
+) -> None:
+    """Analyze kernel fusion differences between AMD and NVIDIA traces.
+
+    Args:
+        trace1: Path to first trace file (AMD or NVIDIA)
+        trace2: Path to second trace file (AMD or NVIDIA)
+        output: Optional output file path (default: stdout)
+        format_type: Output format ('text', 'csv', or 'json')
+        min_group_size: Minimum correlation group size to analyze
+    """
+    # Validate files exist
+    if not trace1.exists():
+        typer.secho(f"❌ File not found: {trace1}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    if not trace2.exists():
+        typer.secho(f"❌ File not found: {trace2}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Analyze fusion
+    # Only show progress messages for non-JSON formats (JSON needs clean stdout)
+    if format_type != 'json':
+        typer.echo("📊 Loading traces...")
+    try:
+        results = analyze_fusion_differences(
+            trace1,
+            trace2,
+            min_group_size=min_group_size,
+        )
+    except Exception as e:
+        typer.secho(
+            f"❌ Error analyzing traces: {e}", fg=typer.colors.RED, err=True
+        )
+        import traceback
+
+        traceback.print_exc()
+        raise typer.Exit(1)
+
+    # Show loading confirmation
+    if format_type != 'json':
+        meta = results["metadata"]
+        # Note: fusion analyzer always uses trace1=AMD, trace2=NVIDIA
+        typer.echo(f"✅ Loaded: {meta['trace1_gpu']} vs {meta['trace2_gpu']}")
+        typer.echo(
+            f"Found {meta['trace1_correlation_groups']} trace1 groups and "
+            f"{meta['trace2_correlation_groups']} trace2 groups with ≥{min_group_size} kernels"
+        )
+        typer.echo(f"✅ Matched {meta['matched_groups']} correlation groups")
+        typer.echo()
+
+    # Generate output
+    if format_type == "text":
+        output_str = format_fusion_text(results)
+    elif format_type == "csv":
+        output_str = format_fusion_csv(results)
+    elif format_type == "json":
+        output_str = format_fusion_json(results)
+    else:
+        typer.secho(f"❌ Unknown format: {format_type}", fg=typer.colors.RED, err=True)
+        raise typer.Exit(1)
+
+    # Write output
+    if output:
+        output.write_text(output_str)
+        typer.secho(f"✅ Report saved to {output}", fg=typer.colors.GREEN)
+    else:
+        typer.echo(output_str)

wafer/wevin_cli.py

@@ -15,6 +15,8 @@ from pathlib import Path
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
     from wafer_core.rollouts import Endpoint, Environment
     from wafer_core.rollouts.dtypes import StreamEvent, ToolCall
     from wafer_core.rollouts.templates import TemplateConfig
@@ -145,21 +147,60 @@ class StreamingChunkFrontend:
         pass
 
 
-def _get_wafer_auth() -> tuple[str | None, str | None]:
+def _make_wafer_token_refresh() -> Callable[[], Awaitable[str | None]]:
+    """Create an async callback that refreshes the wafer proxy token via Supabase."""
+    from .auth import load_credentials, refresh_access_token, save_credentials
+
+    async def _refresh() -> str | None:
+        creds = load_credentials()
+        if not creds or not creds.refresh_token:
+            return None
+        try:
+            new_access, new_refresh = refresh_access_token(creds.refresh_token)
+            save_credentials(new_access, new_refresh, creds.email)
+            return new_access
+        except Exception:
+            return None
+
+    return _refresh
+
+
+def _get_wafer_auth(
+    *, no_proxy: bool = False
+) -> tuple[str | None, str | None, Callable[[], Awaitable[str | None]] | None]:
     """Get wafer auth credentials with fallback chain.
 
     Returns:
-        (api_base, api_key) or (None, None) if no auth found
+        (api_base, api_key, api_key_refresh) or (None, None, None) if no auth found.
+        api_key_refresh is an async callback for mid-session token refresh (only set
+        when using wafer proxy via credentials file).
     """
     from .auth import get_valid_token, load_credentials
     from .global_config import get_api_url
 
+    if no_proxy:
+        api_key = os.environ.get("ANTHROPIC_API_KEY", "")
+        if not api_key:
+            # Try auth.json stored key
+            from wafer_core.auth import get_api_key
+
+            api_key = get_api_key("anthropic") or ""
+        if api_key:
+            print("🔑 Using ANTHROPIC_API_KEY (--no-proxy)\n", file=sys.stderr)
+            return "https://api.anthropic.com", api_key, None
+        print(
+            "❌ --no-proxy requires ANTHROPIC_API_KEY env var or `wafer auth login anthropic`\n",
+            file=sys.stderr,
+        )
+        return None, None, None
+
     # Check WAFER_AUTH_TOKEN env var first
     wafer_token = os.environ.get("WAFER_AUTH_TOKEN", "")
     token_source = "WAFER_AUTH_TOKEN" if wafer_token else None
 
     # Try credentials file with automatic refresh
     had_credentials = False
+    uses_credentials_file = False
     if not wafer_token:
         try:
             creds = load_credentials()
@@ -169,12 +210,16 @@ def _get_wafer_auth() -> tuple[str | None, str | None]:
         wafer_token = get_valid_token()
         if wafer_token:
             token_source = "~/.wafer/credentials.json"
+            uses_credentials_file = True
 
     # If we have a valid wafer token, use it
     if wafer_token:
         api_url = get_api_url()
         print(f"🔑 Using wafer proxy ({token_source})\n", file=sys.stderr)
-        return f"{api_url}/v1/anthropic", wafer_token
+        # Only provide refresh callback when token came from credentials file
+        # (env var tokens are managed externally)
+        refresh = _make_wafer_token_refresh() if uses_credentials_file else None
+        return f"{api_url}/v1/anthropic", wafer_token, refresh
 
     # Fall back to direct anthropic
     api_key = os.environ.get("ANTHROPIC_API_KEY", "")
@@ -186,9 +231,9 @@ def _get_wafer_auth() -> tuple[str | None, str | None]:
             )
         else:
             print("🔑 Using ANTHROPIC_API_KEY\n", file=sys.stderr)
-        return "https://api.anthropic.com", api_key
+        return "https://api.anthropic.com", api_key, None
 
-    return None, None
+    return None, None, None
 
 
 def _get_session_preview(session: object) -> str:
@@ -205,10 +250,22 @@ def _get_session_preview(session: object) -> str:
     return ""
 
 
+def _get_log_file_path() -> Path:
+    """Get user-specific log file path, creating directory if needed.
+
+    Uses ~/.wafer/logs/ to avoid permission issues with shared /tmp.
+    """
+    log_dir = Path.home() / ".wafer" / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+    return log_dir / "wevin_debug.log"
+
+
 def _setup_logging() -> None:
     """Configure logging to file only (no console spam)."""
     import logging.config
 
+    log_file = _get_log_file_path()
+
     logging.config.dictConfig({
         "version": 1,
         "disable_existing_loggers": False,
@@ -220,7 +277,7 @@ def _setup_logging() -> None:
         "handlers": {
             "file": {
                 "class": "logging.handlers.RotatingFileHandler",
-                "filename": "/tmp/wevin_debug.log",
+                "filename": str(log_file),
                 "maxBytes": 10_000_000,
                 "backupCount": 3,
                 "formatter": "json",
@@ -243,6 +300,7 @@ def _build_endpoint(
     model_override: str | None,
     api_base: str,
     api_key: str,
+    api_key_refresh: Callable[[], Awaitable[str | None]] | None = None,
 ) -> Endpoint:
     """Build an Endpoint from template config and auth."""
     from wafer_core.rollouts import Endpoint
@@ -257,6 +315,7 @@ def _build_endpoint(
         model=model_id,
         api_base=api_base,
         api_key=api_key,
+        api_key_refresh=api_key_refresh,
         thinking=thinking_config,
         max_tokens=tpl.max_tokens,
     )
@@ -372,6 +431,7 @@ def main(  # noqa: PLR0913, PLR0915
     get_session: str | None = None,
     json_output: bool = False,
     no_sandbox: bool = False,
+    no_proxy: bool = False,
 ) -> None:
     """Run wevin agent in-process via rollouts."""
     from dataclasses import asdict
@@ -487,10 +547,10 @@ def main(  # noqa: PLR0913, PLR0915
     _setup_logging()
 
     # Auth
-    api_base, api_key = _get_wafer_auth()
+    api_base, api_key, api_key_refresh = _get_wafer_auth(no_proxy=no_proxy)
     if not api_base or not api_key:
         print("Error: No API credentials found", file=sys.stderr)
-        print("  Run 'wafer login' or set ANTHROPIC_API_KEY", file=sys.stderr)
+        print("  Run 'wafer auth login' or set ANTHROPIC_API_KEY", file=sys.stderr)
         sys.exit(1)
 
     assert api_base is not None
@@ -513,6 +573,17 @@ def main(  # noqa: PLR0913, PLR0915
         tpl = _get_default_template()
         base_system_prompt = tpl.system_prompt
 
+    # Compose CLI instructions from --help text for allowed wafer commands
+    # TODO: The eval path doesn't have the skills layer below. If include_skills
+    # is ever enabled for optimize-kernelbench, the eval would need it too for parity.
+    # See test_eval_cli_parity.py for coverage notes.
+    if tpl.bash_allowlist:
+        from wafer.cli_instructions import build_cli_instructions
+
+        cli_instructions = build_cli_instructions(tpl.bash_allowlist)
+        if cli_instructions:
+            base_system_prompt = base_system_prompt + "\n\n" + cli_instructions
+
     # Append skill metadata if skills are enabled
     if tpl.include_skills:
         from wafer_core.rollouts.skills import discover_skills, format_skill_metadata_for_prompt
@@ -530,7 +601,7 @@ def main(  # noqa: PLR0913, PLR0915
     resolved_single_turn = single_turn if single_turn is not None else tpl.single_turn
 
     # Build endpoint and environment
-    endpoint = _build_endpoint(tpl, model, api_base, api_key)
+    endpoint = _build_endpoint(tpl, model, api_base, api_key, api_key_refresh)
     environment = _build_environment(tpl, tools, corpus_path, no_sandbox)
 
     # Session store