zwarm 2.3.5__py3-none-any.whl → 3.6.0__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/config.py

@@ -37,6 +37,7 @@ class ExecutorConfig:
     sandbox: str = "workspace-write"  # read-only | workspace-write | danger-full-access
     timeout: int = 3600
     reasoning_effort: str | None = "high"  # low | medium | high (default to high for compatibility)
+    # Note: web_search is always enabled via .codex/config.toml (set up by `zwarm init`)
 
 
 @dataclass
@@ -59,8 +60,8 @@ class OrchestratorConfig:
     prompt: str | None = None  # path to prompt yaml
     tools: list[str] = field(default_factory=lambda: ["delegate", "converse", "check_session", "end_session", "bash"])
     max_steps: int = 50
+    max_steps_per_turn: int = 60  # Max tool-call steps before returning to user (pilot mode)
     parallel_delegations: int = 4
-    sync_first: bool = True  # prefer sync mode by default
     compaction: CompactionConfig = field(default_factory=CompactionConfig)
 
     # Directory restrictions for agent delegations
@@ -172,8 +173,8 @@ class ZwarmConfig:
                 "prompt": self.orchestrator.prompt,
                 "tools": self.orchestrator.tools,
                 "max_steps": self.orchestrator.max_steps,
+                "max_steps_per_turn": self.orchestrator.max_steps_per_turn,
                 "parallel_delegations": self.orchestrator.parallel_delegations,
-                "sync_first": self.orchestrator.sync_first,
                 "compaction": {
                     "enabled": self.orchestrator.compaction.enabled,
                     "max_tokens": self.orchestrator.compaction.max_tokens,
@@ -195,15 +196,16 @@ class ZwarmConfig:
         }
 
 
-def load_env(path: Path | None = None) -> None:
+def load_env(path: Path | None = None, base_dir: Path | None = None) -> None:
     """Load .env file if it exists."""
     if path is None:
-        path = Path.cwd() / ".env"
+        base = base_dir or Path.cwd()
+        path = base / ".env"
     if path.exists():
         load_dotenv(path)
 
 
-def load_toml_config(path: Path | None = None) -> dict[str, Any]:
+def load_toml_config(path: Path | None = None, base_dir: Path | None = None) -> dict[str, Any]:
     """
     Load config.toml file.
 
@@ -211,11 +213,16 @@ def load_toml_config(path: Path | None = None) -> dict[str, Any]:
     1. Explicit path (if provided)
     2. .zwarm/config.toml (new standard location)
     3. config.toml (legacy location for backwards compat)
+
+    Args:
+        path: Explicit path to config.toml
+        base_dir: Base directory to search in (defaults to cwd)
     """
     if path is None:
+        base = base_dir or Path.cwd()
         # Try new location first
-        new_path = Path.cwd() / ".zwarm" / "config.toml"
-        legacy_path = Path.cwd() / "config.toml"
+        new_path = base / ".zwarm" / "config.toml"
+        legacy_path = base / "config.toml"
         if new_path.exists():
             path = new_path
         elif legacy_path.exists():
@@ -306,6 +313,7 @@ def load_config(
     toml_path: Path | None = None,
     env_path: Path | None = None,
     overrides: list[str] | None = None,
+    working_dir: Path | None = None,
 ) -> ZwarmConfig:
     """
     Load configuration with full precedence chain:
@@ -314,15 +322,24 @@ def load_config(
     3. YAML config file (if provided)
     4. CLI overrides (--set key=value)
     5. Environment variables (for secrets)
+
+    Args:
+        config_path: Path to YAML config file
+        toml_path: Explicit path to config.toml
+        env_path: Explicit path to .env file
+        overrides: CLI overrides (--set key=value)
+        working_dir: Working directory to search for config files (defaults to cwd).
+                    This is important when using --working-dir flag to ensure
+                    config is loaded from the project directory, not invoke directory.
     """
     # Load .env first (for secrets)
-    load_env(env_path)
+    load_env(env_path, base_dir=working_dir)
 
     # Start with defaults
     config_dict: dict[str, Any] = {}
 
     # Layer in config.toml
-    toml_config = load_toml_config(toml_path)
+    toml_config = load_toml_config(toml_path, base_dir=working_dir)
     if toml_config:
         config_dict = deep_merge(config_dict, toml_config)
 

zwarm/core/costs.py

@@ -0,0 +1,71 @@
+"""
+Token cost estimation for LLM models.
+
+This module re-exports from the centralized model registry.
+For adding new models, edit: zwarm/core/registry.py
+
+Backwards-compatible API preserved for existing code.
+"""
+
+from __future__ import annotations
+
+# Re-export everything from registry for backwards compatibility
+from zwarm.core.registry import (
+    ModelInfo,
+    MODELS,
+    resolve_model,
+    get_adapter_for_model,
+    get_default_model,
+    list_models,
+    list_adapters,
+    get_models_help_text,
+    get_models_table_data,
+    estimate_cost,
+    format_cost,
+    estimate_session_cost,
+)
+
+# Backwards compatibility alias
+ModelPricing = ModelInfo
+
+# Legacy aliases for backwards compatibility
+MODEL_PRICING = {m.canonical: m for m in MODELS}
+MODEL_ALIASES = {}
+for m in MODELS:
+    for alias in m.aliases:
+        MODEL_ALIASES[alias] = m.canonical
+
+
+def get_pricing(model: str) -> ModelInfo | None:
+    """
+    Get pricing for a model.
+
+    Args:
+        model: Model name or alias
+
+    Returns:
+        ModelInfo or None if unknown
+    """
+    return resolve_model(model)
+
+
+__all__ = [
+    # New API
+    "ModelInfo",
+    "MODELS",
+    "resolve_model",
+    "get_adapter_for_model",
+    "get_default_model",
+    "list_models",
+    "list_adapters",
+    "get_models_help_text",
+    "get_models_table_data",
+    "estimate_cost",
+    "format_cost",
+    "estimate_session_cost",
+    # Legacy API
+    "MODEL_PRICING",
+    "MODEL_ALIASES",
+    "ModelPricing",
+    "get_pricing",
+]