zwarm 2.3.5__py3-none-any.whl → 3.6.0__py3-none-any.whl

zwarm/core/registry.py

@@ -0,0 +1,329 @@
+"""
+Model Registry - Centralized LLM model definitions for zwarm.
+
+This registry defines all supported models with:
+- Canonical names and aliases
+- Adapter mapping (which CLI handles the model)
+- Pricing information
+
+Add new models here and they'll automatically appear in:
+- `zwarm interactive` help and `models` command
+- Cost estimation
+- Adapter auto-detection from model name
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ModelInfo:
+    """Complete information about an LLM model."""
+
+    # Identity
+    canonical: str  # Full model name (e.g., "gpt-5.1-codex-mini")
+    adapter: str  # "codex" or "claude"
+    aliases: list[str] = field(default_factory=list)  # Short names
+
+    # Pricing ($ per million tokens)
+    input_per_million: float = 0.0
+    output_per_million: float = 0.0
+    cached_input_per_million: float | None = None
+
+    # Metadata
+    description: str = ""
+    is_default: bool = False  # Default model for this adapter
+
+    def estimate_cost(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        cached_tokens: int = 0,
+    ) -> float:
+        """Estimate cost in dollars."""
+        input_cost = (input_tokens / 1_000_000) * self.input_per_million
+        output_cost = (output_tokens / 1_000_000) * self.output_per_million
+
+        cached_cost = 0.0
+        if cached_tokens and self.cached_input_per_million:
+            cached_cost = (cached_tokens / 1_000_000) * self.cached_input_per_million
+
+        return input_cost + output_cost + cached_cost
+
+
+# =============================================================================
+# Model Registry - ADD NEW MODELS HERE
+# =============================================================================
+
+MODELS: list[ModelInfo] = [
+    # -------------------------------------------------------------------------
+    # OpenAI Codex Models (via `codex` CLI)
+    # -------------------------------------------------------------------------
+    ModelInfo(
+        canonical="gpt-5.1-codex-mini",
+        adapter="codex",
+        aliases=["codex-mini", "mini"],
+        input_per_million=0.25,
+        output_per_million=2.00,
+        cached_input_per_million=0.025,
+        description="Fast, cost-effective coding model",
+        is_default=True,
+    ),
+    ModelInfo(
+        canonical="gpt-5.1-codex",
+        adapter="codex",
+        aliases=["codex", "codex-full"],
+        input_per_million=1.25,
+        output_per_million=10.00,
+        cached_input_per_million=0.125,
+        description="Full Codex model with extended reasoning",
+    ),
+    ModelInfo(
+        canonical="gpt-5.1-codex-max",
+        adapter="codex",
+        aliases=["codex-max", "max"],
+        input_per_million=1.25,
+        output_per_million=10.00,
+        cached_input_per_million=0.125,
+        description="Maximum context Codex model",
+    ),
+    # -------------------------------------------------------------------------
+    # Anthropic Claude Models (via `claude` CLI)
+    # -------------------------------------------------------------------------
+    ModelInfo(
+        canonical="sonnet",
+        adapter="claude",
+        aliases=["claude-sonnet", "claude-4-sonnet"],
+        input_per_million=3.00,
+        output_per_million=15.00,
+        description="Balanced Claude model for most tasks",
+        is_default=True,
+    ),
+    ModelInfo(
+        canonical="opus",
+        adapter="claude",
+        aliases=["claude-opus", "claude-4-opus"],
+        input_per_million=15.00,
+        output_per_million=75.00,
+        description="Most capable Claude model",
+    ),
+    ModelInfo(
+        canonical="haiku",
+        adapter="claude",
+        aliases=["claude-haiku", "claude-4-haiku"],
+        input_per_million=0.25,
+        output_per_million=1.25,
+        description="Fast, lightweight Claude model",
+    ),
+]
+
+
+# =============================================================================
+# Registry Lookups
+# =============================================================================
+
+
+def _build_lookup_tables() -> tuple[dict[str, ModelInfo], dict[str, ModelInfo]]:
+    """Build lookup tables for fast model resolution."""
+    by_canonical: dict[str, ModelInfo] = {}
+    by_alias: dict[str, ModelInfo] = {}
+
+    for model in MODELS:
+        by_canonical[model.canonical.lower()] = model
+        by_alias[model.canonical.lower()] = model
+        for alias in model.aliases:
+            by_alias[alias.lower()] = model
+
+    return by_canonical, by_alias
+
+
+_BY_CANONICAL, _BY_ALIAS = _build_lookup_tables()
+
+
+def resolve_model(name: str) -> ModelInfo | None:
+    """
+    Resolve a model name or alias to its ModelInfo.
+
+    Args:
+        name: Model name, alias, or partial match
+
+    Returns:
+        ModelInfo or None if not found
+    """
+    name_lower = name.lower()
+
+    # Exact match on alias or canonical
+    if name_lower in _BY_ALIAS:
+        return _BY_ALIAS[name_lower]
+
+    # Prefix match (e.g., "gpt-5.1-codex-mini-2026-01" -> "gpt-5.1-codex-mini")
+    for canonical, model in _BY_CANONICAL.items():
+        if name_lower.startswith(canonical):
+            return model
+
+    return None
+
+
+def get_adapter_for_model(name: str) -> str | None:
+    """
+    Get the adapter name for a model.
+
+    Args:
+        name: Model name or alias
+
+    Returns:
+        Adapter name ("codex" or "claude") or None if unknown
+    """
+    model = resolve_model(name)
+    return model.adapter if model else None
+
+
+def get_default_model(adapter: str) -> str | None:
+    """
+    Get the default model for an adapter.
+
+    Args:
+        adapter: Adapter name ("codex" or "claude")
+
+    Returns:
+        Default model canonical name or None
+    """
+    for model in MODELS:
+        if model.adapter == adapter and model.is_default:
+            return model.canonical
+    return None
+
+
+def list_models(adapter: str | None = None) -> list[ModelInfo]:
+    """
+    List available models.
+
+    Args:
+        adapter: Filter by adapter, or None for all
+
+    Returns:
+        List of ModelInfo objects
+    """
+    if adapter:
+        return [m for m in MODELS if m.adapter == adapter]
+    return MODELS.copy()
+
+
+def list_adapters() -> list[str]:
+    """Get list of unique adapter names."""
+    return sorted(set(m.adapter for m in MODELS))
+
+
+def get_models_help_text() -> str:
+    """
+    Generate help text showing all available models.
+
+    Returns formatted string for display in help messages.
+    """
+    lines = ["", "Available models:"]
+
+    for adapter in list_adapters():
+        lines.append(f"\n  {adapter.upper()}:")
+        for model in list_models(adapter):
+            default_marker = " *" if model.is_default else ""
+            aliases = ", ".join(model.aliases) if model.aliases else ""
+            alias_str = f" ({aliases})" if aliases else ""
+
+            lines.append(f"    {model.canonical}{alias_str}{default_marker}")
+
+    lines.append("\n  * = default for adapter")
+    return "\n".join(lines)
+
+
+def get_models_table_data() -> list[dict[str, Any]]:
+    """
+    Get model data formatted for table display.
+
+    Returns list of dicts with keys: adapter, model, aliases, default, price, description
+    """
+    data = []
+    for model in MODELS:
+        data.append({
+            "adapter": model.adapter,
+            "model": model.canonical,
+            "aliases": ", ".join(model.aliases),
+            "default": model.is_default,
+            "input_price": model.input_per_million,
+            "output_price": model.output_per_million,
+            "description": model.description,
+        })
+    return data
+
+
+# =============================================================================
+# Cost Estimation
+# =============================================================================
+
+
+def estimate_cost(
+    model: str,
+    input_tokens: int,
+    output_tokens: int,
+    cached_tokens: int = 0,
+) -> float | None:
+    """
+    Estimate cost for a model run.
+
+    Args:
+        model: Model name or alias
+        input_tokens: Number of input tokens
+        output_tokens: Number of output tokens
+        cached_tokens: Number of cached input tokens
+
+    Returns:
+        Cost in USD, or None if model unknown
+    """
+    model_info = resolve_model(model)
+    if model_info is None:
+        return None
+
+    return model_info.estimate_cost(input_tokens, output_tokens, cached_tokens)
+
+
+def format_cost(cost: float | None) -> str:
+    """Format cost as a human-readable string."""
+    if cost is None:
+        return "?"
+    if cost < 0.01:
+        return f"${cost:.4f}"
+    elif cost < 1.00:
+        return f"${cost:.3f}"
+    else:
+        return f"${cost:.2f}"
+
+
+def estimate_session_cost(
+    model: str,
+    token_usage: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Estimate cost for a session given its token usage.
+
+    Args:
+        model: Model used
+        token_usage: Dict with input_tokens, output_tokens, etc.
+
+    Returns:
+        Dict with cost info: {cost, cost_formatted, pricing_known, ...}
+    """
+    input_tokens = token_usage.get("input_tokens", 0)
+    output_tokens = token_usage.get("output_tokens", 0)
+    cached_tokens = token_usage.get("cached_tokens", 0)
+
+    cost = estimate_cost(model, input_tokens, output_tokens, cached_tokens)
+
+    return {
+        "cost": cost,
+        "cost_formatted": format_cost(cost),
+        "pricing_known": cost is not None,
+        "model": model,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+    }

zwarm/core/test_config.py

@@ -20,7 +20,6 @@ def test_default_config():
     assert config.executor.adapter == "codex_mcp"
     assert config.executor.sandbox == "workspace-write"
     assert config.orchestrator.lm == "gpt-5-mini"
-    assert config.orchestrator.sync_first is True
     assert config.state_dir == ".zwarm"
 
 
@@ -68,8 +67,8 @@ def test_apply_overrides():
     assert result["executor"]["adapter"] == "claude_code"
 
     # Override with boolean
-    result = apply_overrides(config, ["orchestrator.sync_first=false"])
-    assert result["orchestrator"]["sync_first"] is False
+    result = apply_overrides(config, ["executor.web_search=true"])
+    assert result["executor"]["web_search"] is True
 
     # Create new nested path
     result = apply_overrides(config, ["weave.project=my-project"])

zwarm/orchestrator.py

@@ -23,7 +23,6 @@ from wbal.helper import TOOL_CALL_TYPE, format_openai_tool_response
 from wbal.lm import LM as wbalLMGeneric
 from wbal.lm import GPT5LargeVerbose
 
-from zwarm.adapters import ExecutorAdapter, get_adapter
 from zwarm.core.compact import compact_messages, should_compact
 from zwarm.core.config import ZwarmConfig, load_config
 from zwarm.core.environment import OrchestratorEnv
@@ -72,7 +71,6 @@ class Orchestrator(YamlAgent):
     # State management
     _state: StateManager = PrivateAttr()
     _sessions: dict[str, ConversationSession] = PrivateAttr(default_factory=dict)
-    _adapters: dict[str, ExecutorAdapter] = PrivateAttr(default_factory=dict)
     _watcher_manager: WatcherManager | None = PrivateAttr(default=None)
     _resumed: bool = PrivateAttr(default=False)
     _total_tokens: int = PrivateAttr(default=0)  # Cumulative orchestrator tokens
@@ -83,9 +81,11 @@ class Orchestrator(YamlAgent):
             "total_tokens": 0,
         }
     )
+    # Callback for step progress (used by CLI to print tool calls)
+    _step_callback: Callable[[int, list[tuple[dict[str, Any], Any]]], None] | None = PrivateAttr(default=None)
 
     def model_post_init(self, __context: Any) -> None:
-        """Initialize state and adapters after model creation."""
+        """Initialize state after model creation."""
         super().model_post_init(__context)
 
         # Initialize state manager with instance isolation
@@ -151,40 +151,9 @@ class Orchestrator(YamlAgent):
         """Access state manager."""
         return self._state
 
-    def _get_adapter(self, name: str) -> ExecutorAdapter:
-        """Get or create an adapter by name using the adapter registry."""
-        if name not in self._adapters:
-            # Get model from config (adapters have their own defaults if None)
-            model = self.config.executor.model
-
-            # Use isolated codex config if available
-            config_path = self.working_dir / self.config.state_dir / "codex.toml"
-            if not config_path.exists():
-                config_path = None  # Fallback to adapter defaults
-
-            self._adapters[name] = get_adapter(
-                name, model=model, config_path=config_path
-            )
-        return self._adapters[name]
-
     def get_executor_usage(self) -> dict[str, int]:
-        """Get aggregated token usage across all executors."""
-        total = {
-            "input_tokens": 0,
-            "output_tokens": 0,
-            "total_tokens": 0,
-        }
-        for adapter in self._adapters.values():
-            if hasattr(adapter, "total_usage"):
-                usage = adapter.total_usage
-                for key in total:
-                    total[key] += usage.get(key, 0)
-        return total
-
-    @property
-    def executor_usage(self) -> dict[str, int]:
-        """Aggregated executor token usage (for Weave tracking)."""
-        return self.get_executor_usage()
+        """Get aggregated token usage from executor sessions."""
+        return self._executor_usage
 
     def save_state(self) -> None:
         """Save orchestrator state for resume."""
@@ -587,7 +556,11 @@ Review what was accomplished in the previous session and delegate new tasks as n
                 }
             # NUDGE and CONTINUE just continue
 
-            self.step()
+            tool_results = self.step()
+
+            # Call step callback if registered (for CLI progress display)
+            if self._step_callback:
+                self._step_callback(self._step_count, tool_results)
 
             if self.stopCondition:
                 break
@@ -599,8 +572,7 @@ Review what was accomplished in the previous session and delegate new tasks as n
 
     async def cleanup(self) -> None:
         """Clean up resources."""
-        for adapter in self._adapters.values():
-            await adapter.cleanup()
+        pass  # Session cleanup handled by CodexSessionManager
 
 
 def build_orchestrator(
@@ -631,15 +603,17 @@ def build_orchestrator(
     """
     from uuid import uuid4
 
-    # Load configuration
+    # Resolve working directory first (needed for config loading)
+    working_dir = working_dir or Path.cwd()
+
+    # Load configuration from working_dir (not cwd!)
+    # This ensures config.toml and .env are loaded from the project being worked on
     config = load_config(
         config_path=config_path,
         overrides=overrides,
+        working_dir=working_dir,
     )
 
-    # Resolve working directory
-    working_dir = working_dir or Path.cwd()
-
     # Generate instance ID if not provided (enables isolation by default for new runs)
     # For resume, instance_id should be provided explicitly
     if instance_id is None and not resume:

zwarm/prompts/__init__.py

@@ -3,8 +3,11 @@ System prompts for zwarm agents.
 """
 
 from zwarm.prompts.orchestrator import ORCHESTRATOR_SYSTEM_PROMPT, get_orchestrator_prompt
+from zwarm.prompts.pilot import PILOT_SYSTEM_PROMPT, get_pilot_prompt
 
 __all__ = [
     "ORCHESTRATOR_SYSTEM_PROMPT",
     "get_orchestrator_prompt",
+    "PILOT_SYSTEM_PROMPT",
+    "get_pilot_prompt",
 ]

zwarm/prompts/orchestrator.py

@@ -27,18 +27,20 @@ For everything else, make your best judgment and proceed. If you're unsure wheth
 
 Your primary tools are for delegation and verification:
 
-**delegate(task, working_dir=None, model=None, wait=True)** - Start a new executor session. The `task` should be a clear, specific description of what you want done. Use `wait=True` (default) for interactive work where you'll iterate with the executor. Use `wait=False` to spawn background work and continue immediately. The `working_dir` parameter lets you run the executor in a specific directory.
+**delegate(task, working_dir=None, model=None)** - Start a new executor session. The `task` should be a clear, specific description of what you want done. All sessions run asynchronously - you'll get a session_id back immediately and can poll for results. The `working_dir` parameter lets you run the executor in a specific directory.
 
-**converse(session_id, message, wait=True)** - Continue an existing conversation. Use this to provide feedback, ask for changes, or guide the executor through complex work. The executor maintains full context. Use `wait=False` to send the message and continue without waiting for a response.
+**converse(session_id, message)** - Continue an existing conversation. Use this to provide feedback, ask for changes, or guide the executor through complex work. The executor maintains full context. Returns immediately - use polling to check for the response.
 
-**peek_session(session_id)** - Quick status check. Returns just the session status and latest message. Use this for fast polling when you have multiple sessions running.
+**peek_session(session_id)** - Quick status check. Returns just the session status and latest message. Use this for fast polling.
 
 **check_session(session_id)** - Full session details including all messages, token usage, runtime. Use this when you need the complete picture.
 
-**list_sessions(status=None)** - List all sessions. Returns a `needs_attention` flag for each session indicating if it recently completed or failed. Use this to monitor multiple parallel sessions and see which ones have new responses ready for review.
+**list_sessions(status=None)** - List all sessions. Returns a `needs_attention` flag for each session indicating if it recently completed or failed. Use this to monitor multiple sessions and see which ones have new responses ready for review.
 
 **end_session(session_id, reason=None, delete=False)** - Kill a running session or clean up a completed one. Use `delete=True` to remove the session entirely (won't show in list_sessions anymore).
 
+**sleep(seconds)** - Pause execution for specified seconds (max 300). Use this when you've started sessions and want to give them time to complete before polling. Essential for the async workflow pattern.
+
 **bash(command)** - Run shell commands directly. Use this primarily for verification: running tests, type checkers, linters, build commands, or inspecting the filesystem. Do NOT use bash to write code yourself - that's what executors are for.
 
 **chat(message, wait_for_user_input)** - Communicate with the human user. Use this sparingly. Most of the time you should be working autonomously without bothering the user.
@@ -63,35 +65,40 @@ The watchers are on your side. They exist to help you succeed, not to criticize.
 
 ---
 
-# Sync vs Async: Choosing the Right Approach
-
-The `wait` parameter controls whether you block waiting for a response or continue immediately.
+# Async Workflow Pattern
 
-**Sync (wait=True)** creates an interactive conversation with the executor. After your task description, you receive the executor's response immediately. You can then provide feedback via converse(), ask for changes, or confirm the work is acceptable. This back-and-forth continues until you're satisfied.
+All executor sessions run asynchronously. When you call delegate() or converse(), you get a session_id back immediately and the executor works in the background. This lets you parallelize work efficiently.
 
-Use sync when the task involves ambiguity, when you expect to iterate, when you want to review results before proceeding, or for high-stakes work needing close supervision.
+The core workflow pattern is: **delegate → sleep → poll → respond**
 
-Typical sync pattern:
-1. `delegate(task)` - get initial response
-2. Evaluate - does it meet requirements?
-3. `converse(id, "feedback...")` - if changes needed
-4. Repeat until satisfied
-5. `end_session(id)` or just move on
+```
+1. delegate(task1) → session_a
+2. delegate(task2) → session_b
+3. delegate(task3) → session_c
+4. sleep(30) → give them time to work
+5. list_sessions() → check which have needs_attention=True
+6. peek_session(a) → quick status check
+7. If still running, sleep(30) and repeat
+8. check_session(a) → full results when done
+9. converse(a, "feedback...") → continue the conversation
+10. sleep(15) → wait for response
+11. check_session(a) → see the response
+```
 
-**Async (wait=False)** is fire-and-forget. You spawn the work and continue immediately without waiting. The executor works in the background.
+**Key principles:**
 
-Use async when tasks are well-defined and self-contained, when you're confident the executor can complete without guidance, or when you want to parallelize multiple independent pieces of work. Async is efficient for clear-cut tasks like "add tests for this function" or "fix this lint error".
+- Use **sleep()** to give executors time to work before polling. Don't spam peek_session() in a tight loop.
+- Use **list_sessions()** to see which sessions have `needs_attention=True` (recently completed or failed).
+- Use **peek_session()** for quick status checks during polling.
+- Use **check_session()** to get full details including all messages when you need to review the actual work.
+- After **converse()**, always sleep() and poll - you won't get the response immediately.
 
-Async pattern for parallel work:
-1. `delegate(task1, wait=False)` → session a
-2. `delegate(task2, wait=False)` → session b
-3. `delegate(task3, wait=False)` → session c
-4. `list_sessions()` → check `needs_attention` flags
-5. `peek_session(a)` → quick status check
-6. `check_session(b)` → full details when ready
-7. `converse(a, "now do X", wait=False)` → continue without blocking
+**Sleep timing guidance:**
 
-When in doubt, prefer sync. The overhead of waiting is small compared to an executor going off in the wrong direction unsupervised.
+- Simple tasks (single file edits, small fixes): 15-30 seconds
+- Medium tasks (multiple files, tests): 30-60 seconds
+- Complex tasks (new features, refactoring): 60-120 seconds
+- If a session is still running after polling, sleep again rather than waiting forever
 
 ---
 
@@ -119,7 +126,7 @@ Never mark work as complete without verifying it actually works. This is the mos
 
 After an executor completes work, run the relevant verification commands. For Python projects, this typically means: pytest for tests, mypy or pyright for type checking, ruff or flake8 for linting. For JavaScript/TypeScript: npm test, tsc for type checking, eslint for linting. For compiled languages: ensure the build succeeds without errors.
 
-When verification fails, you have two options. If you're in a sync session, use converse() to share the error output and ask the executor to fix it. Be specific about what failed - paste the actual error message. If you're in an async session or the sync session has become too confused, end it with verdict="failed" and start a fresh session with a clearer task description that incorporates what you learned.
+When verification fails, use converse() to share the error output and ask the executor to fix it. Be specific about what failed - paste the actual error message. Remember to sleep() and poll for the response. If the session has become too confused or gone too far down the wrong path, end it with verdict="failed" and start a fresh session with a clearer task description that incorporates what you learned.
 
 Do not rationalize failures. If the tests don't pass, the work isn't done. If the type checker complains, the work isn't done. If the linter shows errors, the work isn't done. Your job is to ensure quality, and that means holding firm on verification.
 
@@ -131,7 +138,7 @@ Executors will sometimes fail. They might misunderstand the task, produce buggy
 
 When you notice an executor has gone wrong, first diagnose the problem. What specifically is wrong? Is it a misunderstanding of requirements, a technical error, a missing piece of context? Understanding the root cause helps you correct effectively.
 
-For sync sessions, you can often recover through conversation. Explain what's wrong clearly and specifically. Don't just say "this is wrong" - explain why and what you expected instead. Provide the error messages, the failing test output, or a clear description of the incorrect behavior. Give the executor the information they need to fix the issue.
+You can often recover through conversation using converse(). Explain what's wrong clearly and specifically. Don't just say "this is wrong" - explain why and what you expected instead. Provide the error messages, the failing test output, or a clear description of the incorrect behavior. Give the executor the information they need to fix the issue. Then sleep() and poll for their response.
 
 Sometimes a session becomes too confused or goes too far down the wrong path. In these cases, it's better to cut your losses: call end_session() with verdict="failed" and a summary of what went wrong, then start fresh with a new session that has a better task description informed by what you learned.
 
@@ -145,7 +152,7 @@ Complex tasks often require multiple executor sessions, either in sequence or in
 
 For sequential work with dependencies, complete each session fully before starting the next. Don't leave sessions hanging in an ambiguous state while you start new work. This creates confusion and makes it hard to track what's actually done.
 
-For parallel work on independent tasks, you can start multiple async sessions simultaneously. Use check_session() periodically to monitor progress, and end each session properly when complete. Keep mental track of what's running - don't lose track of sessions.
+For parallel work on independent tasks, start multiple sessions and use the sleep-poll pattern to monitor them. Use list_sessions() to see which have needs_attention=True, check_session() for full details, and end each session properly when complete. Keep mental track of what's running - don't lose track of sessions.
 
 Prioritize completing in-progress work before starting new work. A half-finished feature is worth less than nothing - it's technical debt that will confuse future work. Better to have fewer things fully done than many things partially done.