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

zwarm/core/__init__.py

@@ -0,0 +1,20 @@
+"""Core primitives for zwarm."""
+
+from .checkpoints import Checkpoint, CheckpointManager
+from .costs import (
+    estimate_cost,
+    estimate_session_cost,
+    format_cost,
+    get_pricing,
+    ModelPricing,
+)
+
+__all__ = [
+    "Checkpoint",
+    "CheckpointManager",
+    "estimate_cost",
+    "estimate_session_cost",
+    "format_cost",
+    "get_pricing",
+    "ModelPricing",
+]

zwarm/core/checkpoints.py

@@ -0,0 +1,216 @@
+"""
+Checkpoint primitives for state management.
+
+Provides time-travel capability by recording snapshots of state at key points.
+Used by pilot for turn-by-turn checkpointing, and potentially by other
+interfaces that need state restoration.
+
+Topology reminder:
+    orchestrator → pilot → interactive → CodexSessionManager
+
+These primitives sit at the core layer, usable by any interface above.
+"""
+
+from __future__ import annotations
+
+import copy
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+
+@dataclass
+class Checkpoint:
+    """
+    A snapshot of state at a specific point in time.
+
+    Attributes:
+        checkpoint_id: Unique identifier (e.g., turn number)
+        label: Human-readable label (e.g., "T1", "T2")
+        description: What action led to this state
+        state: The actual state snapshot (deep-copied)
+        timestamp: When checkpoint was created
+        metadata: Optional extra data
+    """
+    checkpoint_id: int
+    label: str
+    description: str
+    state: dict[str, Any]
+    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class CheckpointManager:
+    """
+    Manages checkpoints and time travel.
+
+    Maintains a list of checkpoints and a current position. Supports:
+    - Recording new checkpoints
+    - Jumping to any previous checkpoint
+    - Branching (going back and continuing creates new timeline)
+    - History inspection
+
+    Usage:
+        mgr = CheckpointManager()
+
+        # Record state after each action
+        mgr.record(description="Added auth", state={"messages": [...], ...})
+        mgr.record(description="Fixed bug", state={"messages": [...], ...})
+
+        # Jump back
+        cp = mgr.goto(1)  # Go to first checkpoint
+        restored_state = cp.state
+
+        # Continue from there (branches off)
+        mgr.record(description="Different path", state={...})
+    """
+
+    checkpoints: list[Checkpoint] = field(default_factory=list)
+    current_index: int = -1  # -1 = root (before any checkpoints)
+    next_id: int = 1
+    label_prefix: str = "T"  # Labels will be T1, T2, etc.
+
+    def record(
+        self,
+        description: str,
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+    ) -> Checkpoint:
+        """
+        Record a new checkpoint.
+
+        If not at the end of history (i.e., we've gone back), this creates
+        a branch - future checkpoints are discarded.
+
+        Args:
+            description: What action led to this state
+            state: State to snapshot (will be deep-copied)
+            metadata: Optional extra data
+
+        Returns:
+            The created checkpoint
+        """
+        checkpoint = Checkpoint(
+            checkpoint_id=self.next_id,
+            label=f"{self.label_prefix}{self.next_id}",
+            description=description,
+            state=copy.deepcopy(state),
+            metadata=metadata or {},
+        )
+
+        # If we're not at the end, we're branching - truncate future
+        if self.current_index < len(self.checkpoints) - 1:
+            self.checkpoints = self.checkpoints[:self.current_index + 1]
+
+        self.checkpoints.append(checkpoint)
+        self.current_index = len(self.checkpoints) - 1
+        self.next_id += 1
+
+        return checkpoint
+
+    def goto(self, checkpoint_id: int) -> Checkpoint | None:
+        """
+        Jump to a specific checkpoint.
+
+        Args:
+            checkpoint_id: The checkpoint ID to jump to (0 = root)
+
+        Returns:
+            The checkpoint, or None if not found (or root)
+        """
+        if checkpoint_id == 0:
+            # Root state - before any checkpoints
+            self.current_index = -1
+            return None
+
+        for i, cp in enumerate(self.checkpoints):
+            if cp.checkpoint_id == checkpoint_id:
+                self.current_index = i
+                return cp
+
+        return None  # Not found
+
+    def goto_label(self, label: str) -> Checkpoint | None:
+        """
+        Jump to a checkpoint by label (e.g., "T1", "root").
+
+        Args:
+            label: The label to find
+
+        Returns:
+            The checkpoint, or None if not found
+        """
+        if label.lower() == "root":
+            self.current_index = -1
+            return None
+
+        for i, cp in enumerate(self.checkpoints):
+            if cp.label == label:
+                self.current_index = i
+                return cp
+
+        return None
+
+    def current(self) -> Checkpoint | None:
+        """Get the current checkpoint, or None if at root."""
+        if self.current_index < 0 or self.current_index >= len(self.checkpoints):
+            return None
+        return self.checkpoints[self.current_index]
+
+    def current_state(self) -> dict[str, Any] | None:
+        """Get the current state, or None if at root."""
+        cp = self.current()
+        return copy.deepcopy(cp.state) if cp else None
+
+    def history(
+        self,
+        limit: int | None = None,
+        include_state: bool = False,
+    ) -> list[dict[str, Any]]:
+        """
+        Get history entries for display.
+
+        Args:
+            limit: Max entries to return (most recent)
+            include_state: Whether to include full state in entries
+
+        Returns:
+            List of history entries with checkpoint info
+        """
+        entries = []
+        for i, cp in enumerate(self.checkpoints):
+            entry = {
+                "checkpoint_id": cp.checkpoint_id,
+                "label": cp.label,
+                "description": cp.description,
+                "timestamp": cp.timestamp,
+                "is_current": i == self.current_index,
+                "metadata": cp.metadata,
+            }
+            if include_state:
+                entry["state"] = cp.state
+            entries.append(entry)
+
+        if limit:
+            entries = entries[-limit:]
+
+        return entries
+
+    def label_for(self, checkpoint_id: int) -> str:
+        """Get label for a checkpoint ID."""
+        if checkpoint_id == 0:
+            return "root"
+        return f"{self.label_prefix}{checkpoint_id}"
+
+    def __len__(self) -> int:
+        """Number of checkpoints."""
+        return len(self.checkpoints)
+
+    def is_at_root(self) -> bool:
+        """Whether we're at root (before any checkpoints)."""
+        return self.current_index < 0
+
+    def is_at_end(self) -> bool:
+        """Whether we're at the most recent checkpoint."""
+        return self.current_index == len(self.checkpoints) - 1

zwarm/core/costs.py

@@ -0,0 +1,199 @@
+"""
+Token cost estimation for LLM models.
+
+Pricing data is hardcoded and may become stale. Last updated: 2026-01.
+
+Sources:
+- https://www.helicone.ai/llm-cost/provider/openai/model/gpt-5.1-codex
+- https://pricepertoken.com/pricing-page/model/openai-codex-mini
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class ModelPricing:
+    """Pricing for a model in $ per million tokens."""
+    input_per_million: float
+    output_per_million: float
+    cached_input_per_million: float | None = None  # Some models have cached input discount
+
+    def estimate_cost(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        cached_tokens: int = 0,
+    ) -> float:
+        """
+        Estimate cost in dollars.
+
+        Args:
+            input_tokens: Number of input tokens
+            output_tokens: Number of output tokens
+            cached_tokens: Number of cached input tokens (if applicable)
+
+        Returns:
+            Estimated cost in USD
+        """
+        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 pricing table ($ per million tokens)
+# Last updated: 2026-01
+MODEL_PRICING: dict[str, ModelPricing] = {
+    # OpenAI Codex models
+    "gpt-5.1-codex": ModelPricing(
+        input_per_million=1.25,
+        output_per_million=10.00,
+        cached_input_per_million=0.125,  # 90% discount for cached
+    ),
+    "gpt-5.1-codex-mini": ModelPricing(
+        input_per_million=0.25,
+        output_per_million=2.00,
+        cached_input_per_million=0.025,
+    ),
+    "gpt-5.1-codex-max": ModelPricing(
+        input_per_million=1.25,
+        output_per_million=10.00,
+        cached_input_per_million=0.125,
+    ),
+    # GPT-5 base models (for reference)
+    "gpt-5": ModelPricing(
+        input_per_million=1.25,
+        output_per_million=10.00,
+    ),
+    "gpt-5-mini": ModelPricing(
+        input_per_million=0.25,
+        output_per_million=2.00,
+    ),
+    # Claude models (Anthropic)
+    "claude-sonnet-4-20250514": ModelPricing(
+        input_per_million=3.00,
+        output_per_million=15.00,
+    ),
+    "claude-opus-4-20250514": ModelPricing(
+        input_per_million=15.00,
+        output_per_million=75.00,
+    ),
+    "claude-3-5-sonnet-20241022": ModelPricing(
+        input_per_million=3.00,
+        output_per_million=15.00,
+    ),
+}
+
+# Aliases for common model names
+MODEL_ALIASES: dict[str, str] = {
+    "codex": "gpt-5.1-codex",
+    "codex-mini": "gpt-5.1-codex-mini",
+    "codex-max": "gpt-5.1-codex-max",
+    "gpt5": "gpt-5",
+    "gpt5-mini": "gpt-5-mini",
+    "sonnet": "claude-sonnet-4-20250514",
+    "opus": "claude-opus-4-20250514",
+}
+
+
+def get_pricing(model: str) -> ModelPricing | None:
+    """
+    Get pricing for a model.
+
+    Args:
+        model: Model name or alias
+
+    Returns:
+        ModelPricing or None if unknown
+    """
+    # Check aliases first
+    resolved = MODEL_ALIASES.get(model.lower(), model)
+
+    # Exact match
+    if resolved in MODEL_PRICING:
+        return MODEL_PRICING[resolved]
+
+    # Try lowercase
+    if resolved.lower() in MODEL_PRICING:
+        return MODEL_PRICING[resolved.lower()]
+
+    # Try prefix matching (e.g., "gpt-5.1-codex-mini-2026-01" -> "gpt-5.1-codex-mini")
+    for known_model in MODEL_PRICING:
+        if resolved.lower().startswith(known_model.lower()):
+            return MODEL_PRICING[known_model]
+
+    return None
+
+
+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
+        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 pricing unknown
+    """
+    pricing = get_pricing(model)
+    if pricing is None:
+        return None
+
+    return pricing.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/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.