wafer-cli 0.2.14__py3-none-any.whl → 0.2.30__py3-none-any.whl

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,
     )
@@ -266,18 +325,27 @@ def _build_environment(
     tpl: TemplateConfig,
     tools_override: list[str] | None,
     corpus_path: str | None,
+    no_sandbox: bool = False,
 ) -> Environment:
     """Build a CodingEnvironment from template config."""
     from wafer_core.environments.coding import CodingEnvironment
     from wafer_core.rollouts.templates import DANGEROUS_BASH_COMMANDS
+    from wafer_core.sandbox import SandboxMode
 
     working_dir = Path(corpus_path) if corpus_path else Path.cwd()
-    resolved_tools = tools_override or tpl.tools
+    resolved_tools = list(tools_override or tpl.tools)
+
+    # Add skill tool if skills are enabled
+    if tpl.include_skills and "skill" not in resolved_tools:
+        resolved_tools.append("skill")
+
+    sandbox_mode = SandboxMode.DISABLED if no_sandbox else SandboxMode.ENABLED
     env: Environment = CodingEnvironment(
         working_dir=working_dir,
         enabled_tools=resolved_tools,
         bash_allowlist=tpl.bash_allowlist,
         bash_denylist=DANGEROUS_BASH_COMMANDS,
+        sandbox_mode=sandbox_mode,
     )  # type: ignore[assignment]
     return env
 
@@ -362,6 +430,8 @@ def main(  # noqa: PLR0913, PLR0915
     list_sessions: bool = False,
     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
@@ -373,6 +443,7 @@ def main(  # noqa: PLR0913, PLR0915
 
     # Handle --get-session: load session by ID and print
     if get_session:
+
         async def _get_session() -> None:
             try:
                 session, err = await session_store.get(get_session)
@@ -393,16 +464,18 @@ def main(  # noqa: PLR0913, PLR0915
                         error_msg = f"Failed to serialize messages: {e}"
                         print(json.dumps({"error": error_msg}))
                         sys.exit(1)
-                    
-                    print(json.dumps({
-                        "session_id": session.session_id,
-                        "status": session.status.value,
-                        "model": session.endpoint.model if session.endpoint else None,
-                        "created_at": session.created_at,
-                        "updated_at": session.updated_at,
-                        "messages": messages_data,
-                        "tags": session.tags,
-                    }))
+
+                    print(
+                        json.dumps({
+                            "session_id": session.session_id,
+                            "status": session.status.value,
+                            "model": session.endpoint.model if session.endpoint else None,
+                            "created_at": session.created_at,
+                            "updated_at": session.updated_at,
+                            "messages": messages_data,
+                            "tags": session.tags,
+                        })
+                    )
                 else:
                     print(f"Session: {session.session_id}")
                     print(f"Status: {session.status.value}")
@@ -474,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
@@ -490,7 +563,7 @@ def main(  # noqa: PLR0913, PLR0915
             print(f"Error loading template: {err}", file=sys.stderr)
             sys.exit(1)
         tpl = loaded_template
-        system_prompt = tpl.interpolate_prompt(template_args or {})
+        base_system_prompt = tpl.interpolate_prompt(template_args or {})
         # Show template info when starting without a prompt
         if not prompt and tpl.description:
             print(f"Template: {tpl.name}", file=sys.stderr)
@@ -498,14 +571,38 @@ def main(  # noqa: PLR0913, PLR0915
             print(file=sys.stderr)
     else:
         tpl = _get_default_template()
-        system_prompt = tpl.system_prompt
+        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
+
+        skill_metadata = discover_skills()
+        if skill_metadata:
+            skill_section = format_skill_metadata_for_prompt(skill_metadata)
+            system_prompt = base_system_prompt + "\n\n" + skill_section
+        else:
+            system_prompt = base_system_prompt
+    else:
+        system_prompt = base_system_prompt
 
     # CLI args override template values
     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)
-    environment = _build_environment(tpl, tools, corpus_path)
+    endpoint = _build_endpoint(tpl, model, api_base, api_key, api_key_refresh)
+    environment = _build_environment(tpl, tools, corpus_path, no_sandbox)
 
     # Session store
     session_store = FileSessionStore()
@@ -545,7 +642,7 @@ def main(  # noqa: PLR0913, PLR0915
             else:
                 if json_output:
                     # Emit session_start if we have a session_id (from --resume)
-                    model_name = endpoint.model if hasattr(endpoint, 'model') else None
+                    model_name = endpoint.model if hasattr(endpoint, "model") else None
                     frontend = StreamingChunkFrontend(session_id=session_id, model=model_name)
                 else:
                     frontend = NoneFrontend(show_tool_calls=True, show_thinking=False)
@@ -560,9 +657,11 @@ def main(  # noqa: PLR0913, PLR0915
                 # Emit session_start for new sessions (if session_id was None and we got one)
                 # Check first state to emit as early as possible
                 if json_output and isinstance(frontend, StreamingChunkFrontend):
-                    first_session_id = states[0].session_id if states and states[0].session_id else None
+                    first_session_id = (
+                        states[0].session_id if states and states[0].session_id else None
+                    )
                     if first_session_id and not session_id:  # New session created
-                        model_name = endpoint.model if hasattr(endpoint, 'model') else None
+                        model_name = endpoint.model if hasattr(endpoint, "model") else None
                         frontend.emit_session_start(first_session_id, model_name)
                 # Print resume command with full wafer agent prefix
                 if states and states[-1].session_id:

wafer/workspaces.py

@@ -13,7 +13,7 @@ import httpx
 from .api_client import get_api_url
 from .auth import get_auth_headers
 
-VALID_STATUSES = {"creating", "running"}
+VALID_STATUSES = {"creating", "running", "error"}
 
 
 def _get_client() -> tuple[str, dict[str, str]]:
@@ -39,13 +39,13 @@ def _friendly_error(status_code: int, response_text: str, workspace_id: str) ->
         User-friendly error message with suggested next steps
     """
     if status_code == 401:
-        return "Not authenticated. Run: wafer login"
+        return "Not authenticated. Run: wafer auth login"
 
     if status_code == 402:
         return (
             "Insufficient credits.\n"
-            "  Check usage: wafer billing\n"
-            "  Add credits: wafer billing topup"
+            "  Check usage: wafer config billing\n"
+            "  Add credits: wafer config billing topup"
         )
 
     if status_code == 404:
@@ -107,7 +107,7 @@ def _list_workspaces_raw() -> list[dict]:
             workspaces = response.json()
     except httpx.HTTPStatusError as e:
         if e.response.status_code == 401:
-            raise RuntimeError("Not authenticated. Run: wafer login") from e
+            raise RuntimeError("Not authenticated. Run: wafer auth login") from e
         raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
     except httpx.RequestError as e:
         raise RuntimeError(f"Could not reach API: {e}") from e
@@ -188,7 +188,7 @@ def list_workspaces(json_output: bool = False) -> str:
             workspaces = response.json()
     except httpx.HTTPStatusError as e:
         if e.response.status_code == 401:
-            raise RuntimeError("Not authenticated. Run: wafer login") from e
+            raise RuntimeError("Not authenticated. Run: wafer auth login") from e
         raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
     except httpx.RequestError as e:
         raise RuntimeError(f"Could not reach API: {e}") from e
@@ -211,17 +211,39 @@ def list_workspaces(json_output: bool = False) -> str:
     lines = ["Workspaces:", ""]
     for ws in workspaces:
         status = ws.get("status", "unknown")
-        status_icon = {"running": "●", "creating": "◐"}.get(status, "?")
+        status_icon = {"running": "●", "creating": "◐", "error": "✗"}.get(status, "?")
         lines.append(f"  {status_icon} {ws['name']} ({ws['id']})")
         lines.append(f"    GPU: {ws.get('gpu_type', 'N/A')} | Image: {ws.get('image', 'N/A')}")
-        if ws.get("ssh_host") and ws.get("ssh_port") and ws.get("ssh_user"):
+
+        if status == "error":
+            lines.append(
+                f"    Status: Provisioning failed. Delete and recreate: wafer workspaces delete {ws['name']}"
+            )
+        elif ws.get("ssh_host") and ws.get("ssh_port") and ws.get("ssh_user"):
+            ssh_line = f"    SSH: ssh -p {ws['ssh_port']} {ws['ssh_user']}@{ws['ssh_host']}"
+            if status == "creating":
+                ssh_line += " (finalizing...)"
+            lines.append(ssh_line)
+        elif status == "running":
             lines.append(
-                f"    SSH: ssh -p {ws['ssh_port']} {ws['ssh_user']}@{ws['ssh_host']}"
+                f"    Status: Running but SSH not ready. Try: wafer workspaces delete {ws['name']} && wafer workspaces create {ws['name']} --wait"
             )
         else:
-            lines.append("    SSH: Not ready (run: wafer workspaces ssh <name>)")
+            lines.append("    SSH: Not ready (workspace is still creating)")
         lines.append("")
 
+    # Add SSH tip for users with running workspaces
+    has_running_with_ssh = any(
+        ws.get("status") == "running" and ws.get("ssh_host")
+        for ws in workspaces
+    )
+    if has_running_with_ssh:
+        lines.append("Tip: SSH directly for interactive work. 'exec' is for quick commands only.")
+    
+    has_error = any(ws.get("status") == "error" for ws in workspaces)
+    if has_error:
+        lines.append("Note: Error workspaces are auto-cleaned after 12 hours.")
+
     return "\n".join(lines)
 
 
@@ -285,7 +307,7 @@ def create_workspace(
             workspace = response.json()
     except httpx.HTTPStatusError as e:
         if e.response.status_code == 401:
-            raise RuntimeError("Not authenticated. Run: wafer login") from e
+            raise RuntimeError("Not authenticated. Run: wafer auth login") from e
         if e.response.status_code == 400:
             raise RuntimeError(f"Bad request: {e.response.text}") from e
         raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
@@ -391,7 +413,7 @@ def delete_workspace(workspace_id: str, json_output: bool = False) -> str:
             result = response.json()
     except httpx.HTTPStatusError as e:
         if e.response.status_code == 401:
-            raise RuntimeError("Not authenticated. Run: wafer login") from e
+            raise RuntimeError("Not authenticated. Run: wafer auth login") from e
         if e.response.status_code == 404:
             raise RuntimeError(f"Workspace not found: {workspace_id}") from e
         raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
@@ -441,6 +463,12 @@ def sync_files(
         f"Workspace {workspace_id} has invalid status '{workspace_status}'. "
         f"Valid statuses: {VALID_STATUSES}"
     )
+    if workspace_status == "error":
+        raise RuntimeError(
+            f"Workspace provisioning failed. Delete and recreate:\n"
+            f"  wafer workspaces delete {workspace_id}\n"
+            f"  wafer workspaces create {ws.get('name', workspace_id)} --wait"
+        )
     if workspace_status != "running":
         raise RuntimeError(
             f"Workspace is {workspace_status}. Wait for it to be running before syncing."
@@ -448,9 +476,14 @@ def sync_files(
     ssh_host = ws.get("ssh_host")
     ssh_port = ws.get("ssh_port")
     ssh_user = ws.get("ssh_user")
-    assert ssh_host, "Workspace missing ssh_host"
+    if not ssh_host or not ssh_port or not ssh_user:
+        # Workspace is running but SSH credentials are missing - unusual state
+        raise RuntimeError(
+            f"Workspace is running but SSH not ready.\n"
+            f"  Delete and recreate: wafer workspaces delete {workspace_id}\n"
+            f"  Then: wafer workspaces create {ws.get('name', workspace_id)} --wait"
+        )
     assert isinstance(ssh_port, int) and ssh_port > 0, "Workspace missing valid ssh_port"
-    assert ssh_user, "Workspace missing ssh_user"
 
     # Build rsync command
     # -a: archive mode (preserves permissions, etc.)
@@ -509,6 +542,102 @@ def sync_files(
     return file_count, warning
 
 
+def pull_files(
+    workspace_id: str,
+    remote_path: str,
+    local_path: Path,
+    on_progress: Callable[[str], None] | None = None,
+) -> int:
+    """Pull files from workspace to local via rsync over SSH.
+
+    Args:
+        workspace_id: Workspace ID or name
+        remote_path: Remote path in workspace (relative to /workspace or absolute)
+        local_path: Local destination path
+        on_progress: Optional callback for progress messages
+
+    Returns:
+        Number of files transferred
+
+    Raises:
+        RuntimeError: If rsync fails or workspace not accessible
+    """
+    import subprocess
+
+    def emit(msg: str) -> None:
+        if on_progress:
+            on_progress(msg)
+
+    assert workspace_id, "Workspace ID must be non-empty"
+
+    ws = get_workspace_raw(workspace_id)
+    workspace_status = ws.get("status")
+    assert workspace_status in VALID_STATUSES, (
+        f"Workspace {workspace_id} has invalid status '{workspace_status}'. "
+        f"Valid statuses: {VALID_STATUSES}"
+    )
+    if workspace_status == "error":
+        raise RuntimeError(
+            f"Workspace provisioning failed. Delete and recreate:\n"
+            f"  wafer workspaces delete {workspace_id}\n"
+            f"  wafer workspaces create {ws.get('name', workspace_id)} --wait"
+        )
+    if workspace_status != "running":
+        raise RuntimeError(
+            f"Workspace is {workspace_status}. Wait for it to be running before pulling files."
+        )
+    ssh_host = ws.get("ssh_host")
+    ssh_port = ws.get("ssh_port")
+    ssh_user = ws.get("ssh_user")
+    if not ssh_host or not ssh_port or not ssh_user:
+        raise RuntimeError(
+            f"Workspace is running but SSH not ready.\n"
+            f"  Delete and recreate: wafer workspaces delete {workspace_id}\n"
+            f"  Then: wafer workspaces create {ws.get('name', workspace_id)} --wait"
+        )
+    assert isinstance(ssh_port, int) and ssh_port > 0, "Workspace missing valid ssh_port"
+
+    # Normalize remote path - if not absolute, assume relative to /workspace
+    if not remote_path.startswith("/"):
+        remote_path = f"/workspace/{remote_path}"
+
+    # Build SSH command for rsync
+    ssh_opts = f"-p {ssh_port} -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"
+
+    # Build rsync command (reverse of sync - from remote to local)
+    rsync_cmd = [
+        "rsync",
+        "-avz",
+        "-e",
+        f"ssh {ssh_opts}",
+        f"{ssh_user}@{ssh_host}:{remote_path}",
+        str(local_path),
+    ]
+
+    emit(f"Pulling {remote_path} from workspace...")
+
+    try:
+        result = subprocess.run(rsync_cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(f"rsync failed: {result.stderr}")
+
+        # Count files from rsync output
+        lines = result.stdout.strip().split("\n")
+        file_count = sum(
+            1
+            for line in lines
+            if line and not line.startswith((" ", "sent", "total", "receiving", "building"))
+        )
+
+    except FileNotFoundError:
+        raise RuntimeError("rsync not found. Install rsync to use pull feature.") from None
+    except subprocess.SubprocessError as e:
+        raise RuntimeError(f"Pull failed: {e}") from e
+
+    emit(f"Pulled {file_count} files")
+    return file_count
+
+
 def _init_sync_state(workspace_id: str) -> str | None:
     """Tell API to sync files from bare metal to Modal volume.
 
@@ -562,7 +691,7 @@ def get_workspace_raw(workspace_id: str) -> dict:
             workspace = response.json()
     except httpx.HTTPStatusError as e:
         if e.response.status_code == 401:
-            raise RuntimeError("Not authenticated. Run: wafer login") from e
+            raise RuntimeError("Not authenticated. Run: wafer auth login") from e
         if e.response.status_code == 404:
             raise RuntimeError(f"Workspace not found: {workspace_id}") from e
         raise RuntimeError(f"API error: {e.response.status_code} - {e.response.text}") from e
@@ -607,16 +736,34 @@ def get_workspace(workspace_id: str, json_output: bool = False) -> str:
         f"  Last Used: {workspace.get('last_used_at', 'N/A')}",
     ]
 
-    if workspace.get("ssh_host"):
+    if status == "error":
+        lines.extend([
+            "",
+            "Provisioning failed. Delete and recreate:",
+            f"  wafer workspaces delete {workspace['name']}",
+            f"  wafer workspaces create {workspace['name']} --wait",
+            "",
+            "Note: Error workspaces are auto-cleaned after 12 hours.",
+        ])
+    elif workspace.get("ssh_host"):
         lines.extend([
             "",
             "SSH Info:",
             f"  Host: {workspace['ssh_host']}",
             f"  Port: {workspace.get('ssh_port', 22)}",
             f"  User: {workspace.get('ssh_user', 'root')}",
+            "",
+            "Tip: SSH directly for interactive work. 'exec' is for quick commands only.",
         ])
     elif status == "creating":
         lines.extend(["", "SSH: available once workspace is running"])
+    elif status == "running":
+        # Running but no SSH credentials - unusual state
+        lines.extend([
+            "",
+            "Status: Running but SSH not ready.",
+            f"  Delete and recreate: wafer workspaces delete {workspace['name']}",
+        ])
 
     return "\n".join(lines)