zwarm 0.1.0__py3-none-any.whl

zwarm/core/state.py

@@ -0,0 +1,224 @@
+"""
+Flat-file state management for zwarm.
+
+State structure:
+.zwarm/
+├── state.json              # Current state (sessions, tasks)
+├── events.jsonl            # Append-only event log
+├── sessions/
+│   └── <session-id>/
+│       ├── messages.json   # Full conversation history
+│       └── output.log      # Agent stdout/stderr
+└── orchestrator/
+    └── messages.json       # Orchestrator's message history (for resume)
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from .models import ConversationSession, Event, Task
+
+
+def _json_serializer(obj: Any) -> Any:
+    """Custom JSON serializer for non-standard types."""
+    # Handle pydantic models
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump()
+    # Handle objects with __dict__
+    if hasattr(obj, "__dict__"):
+        return {k: v for k, v in obj.__dict__.items() if not k.startswith("_")}
+    # Handle datetime
+    if hasattr(obj, "isoformat"):
+        return obj.isoformat()
+    # Fallback to string representation
+    return str(obj)
+
+
+class StateManager:
+    """
+    Manages flat-file state for zwarm.
+
+    All state is stored as JSON files in a directory (default: .zwarm/).
+    This enables:
+    - Git-backed history
+    - Easy debugging (just read the files)
+    - Resume from previous state
+    """
+
+    def __init__(self, state_dir: Path | str = ".zwarm"):
+        self.state_dir = Path(state_dir)
+        self._sessions: dict[str, ConversationSession] = {}
+        self._tasks: dict[str, Task] = {}
+        self._orchestrator_messages: list[dict[str, Any]] = []
+
+    def init(self) -> None:
+        """Initialize state directory structure."""
+        self.state_dir.mkdir(parents=True, exist_ok=True)
+        (self.state_dir / "sessions").mkdir(exist_ok=True)
+        (self.state_dir / "orchestrator").mkdir(exist_ok=True)
+
+        # Touch events.jsonl
+        events_file = self.state_dir / "events.jsonl"
+        if not events_file.exists():
+            events_file.touch()
+
+    # --- Sessions ---
+
+    def add_session(self, session: ConversationSession) -> None:
+        """Add a session and persist it."""
+        self._sessions[session.id] = session
+        self._save_session(session)
+        self._save_state()
+
+    def get_session(self, session_id: str) -> ConversationSession | None:
+        """Get a session by ID."""
+        return self._sessions.get(session_id)
+
+    def update_session(self, session: ConversationSession) -> None:
+        """Update a session and persist it."""
+        self._sessions[session.id] = session
+        self._save_session(session)
+        self._save_state()
+
+    def list_sessions(self, status: str | None = None) -> list[ConversationSession]:
+        """List sessions, optionally filtered by status."""
+        sessions = list(self._sessions.values())
+        if status:
+            sessions = [s for s in sessions if s.status.value == status]
+        return sessions
+
+    def _save_session(self, session: ConversationSession) -> None:
+        """Save session to its own directory."""
+        session_dir = self.state_dir / "sessions" / session.id
+        session_dir.mkdir(parents=True, exist_ok=True)
+
+        # Save messages
+        messages_file = session_dir / "messages.json"
+        messages_file.write_text(json.dumps([m.to_dict() for m in session.messages], indent=2))
+
+    # --- Tasks ---
+
+    def add_task(self, task: Task) -> None:
+        """Add a task and persist it."""
+        self._tasks[task.id] = task
+        self._save_state()
+
+    def get_task(self, task_id: str) -> Task | None:
+        """Get a task by ID."""
+        return self._tasks.get(task_id)
+
+    def update_task(self, task: Task) -> None:
+        """Update a task and persist it."""
+        self._tasks[task.id] = task
+        self._save_state()
+
+    def list_tasks(self, status: str | None = None) -> list[Task]:
+        """List tasks, optionally filtered by status."""
+        tasks = list(self._tasks.values())
+        if status:
+            tasks = [t for t in tasks if t.status.value == status]
+        return tasks
+
+    # --- Events ---
+
+    def log_event(self, event: Event) -> None:
+        """Append an event to the log."""
+        events_file = self.state_dir / "events.jsonl"
+        with open(events_file, "a") as f:
+            f.write(json.dumps(event.to_dict()) + "\n")
+
+    def get_events(
+        self,
+        session_id: str | None = None,
+        task_id: str | None = None,
+        kind: str | None = None,
+        limit: int | None = None,
+    ) -> list[Event]:
+        """Read events from the log, optionally filtered."""
+        events_file = self.state_dir / "events.jsonl"
+        if not events_file.exists():
+            return []
+
+        events = []
+        with open(events_file) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                event = Event.from_dict(json.loads(line))
+                if session_id and event.session_id != session_id:
+                    continue
+                if task_id and event.task_id != task_id:
+                    continue
+                if kind and event.kind != kind:
+                    continue
+                events.append(event)
+
+        # Most recent first
+        events.reverse()
+        if limit:
+            events = events[:limit]
+        return events
+
+    # --- Orchestrator State ---
+
+    def save_orchestrator_messages(self, messages: list[dict[str, Any]]) -> None:
+        """Save orchestrator's message history for resume."""
+        self._orchestrator_messages = messages
+        messages_file = self.state_dir / "orchestrator" / "messages.json"
+        # Use custom encoder to handle non-serializable types
+        messages_file.write_text(json.dumps(messages, indent=2, default=_json_serializer))
+
+    def load_orchestrator_messages(self) -> list[dict[str, Any]]:
+        """Load orchestrator's message history for resume."""
+        messages_file = self.state_dir / "orchestrator" / "messages.json"
+        if not messages_file.exists():
+            return []
+        return json.loads(messages_file.read_text())
+
+    # --- State Persistence ---
+
+    def _save_state(self) -> None:
+        """Save current state to state.json."""
+        state = {
+            "updated_at": datetime.now().isoformat(),
+            "sessions": {sid: s.to_dict() for sid, s in self._sessions.items()},
+            "tasks": {tid: t.to_dict() for tid, t in self._tasks.items()},
+        }
+        state_file = self.state_dir / "state.json"
+        state_file.write_text(json.dumps(state, indent=2))
+
+    def load(self) -> None:
+        """Load state from state.json."""
+        state_file = self.state_dir / "state.json"
+        if not state_file.exists():
+            return
+
+        state = json.loads(state_file.read_text())
+
+        # Load sessions
+        for sid, sdata in state.get("sessions", {}).items():
+            self._sessions[sid] = ConversationSession.from_dict(sdata)
+
+        # Load tasks
+        for tid, tdata in state.get("tasks", {}).items():
+            self._tasks[tid] = Task.from_dict(tdata)
+
+    def clear(self) -> None:
+        """Clear all state (for testing)."""
+        self._sessions.clear()
+        self._tasks.clear()
+        self._orchestrator_messages.clear()
+
+        # Clear files
+        state_file = self.state_dir / "state.json"
+        if state_file.exists():
+            state_file.unlink()
+
+        events_file = self.state_dir / "events.jsonl"
+        if events_file.exists():
+            events_file.write_text("")

zwarm/core/test_config.py

@@ -0,0 +1,160 @@
+"""Tests for the config system."""
+
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from zwarm.core.config import (
+    ZwarmConfig,
+    apply_overrides,
+    deep_merge,
+    load_config,
+    load_yaml_config,
+)
+
+
+def test_default_config():
+    """Test default configuration values."""
+    config = ZwarmConfig()
+    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"
+
+
+def test_config_from_dict():
+    """Test creating config from dictionary."""
+    data = {
+        "executor": {"adapter": "claude_code", "model": "claude-sonnet"},
+        "orchestrator": {"lm": "gpt-5", "max_steps": 100},
+        "state_dir": ".my_state",
+    }
+    config = ZwarmConfig.from_dict(data)
+    assert config.executor.adapter == "claude_code"
+    assert config.executor.model == "claude-sonnet"
+    assert config.orchestrator.lm == "gpt-5"
+    assert config.orchestrator.max_steps == 100
+    assert config.state_dir == ".my_state"
+
+
+def test_config_to_dict():
+    """Test converting config to dictionary."""
+    config = ZwarmConfig()
+    data = config.to_dict()
+    assert data["executor"]["adapter"] == "codex_mcp"
+    assert data["orchestrator"]["lm"] == "gpt-5-mini"
+
+
+def test_deep_merge():
+    """Test deep merging of dictionaries."""
+    base = {"a": 1, "b": {"c": 2, "d": 3}}
+    override = {"b": {"c": 99}, "e": 4}
+    result = deep_merge(base, override)
+    assert result == {"a": 1, "b": {"c": 99, "d": 3}, "e": 4}
+
+
+def test_apply_overrides():
+    """Test CLI override application."""
+    config = {"executor": {"adapter": "codex"}, "orchestrator": {"max_steps": 10}}
+
+    # Override nested value
+    result = apply_overrides(config, ["orchestrator.max_steps=50"])
+    assert result["orchestrator"]["max_steps"] == 50
+
+    # Override with string
+    result = apply_overrides(config, ["executor.adapter=claude_code"])
+    assert result["executor"]["adapter"] == "claude_code"
+
+    # Override with boolean
+    result = apply_overrides(config, ["orchestrator.sync_first=false"])
+    assert result["orchestrator"]["sync_first"] is False
+
+    # Create new nested path
+    result = apply_overrides(config, ["weave.project=my-project"])
+    assert result["weave"]["project"] == "my-project"
+
+
+def test_yaml_inheritance():
+    """Test YAML config inheritance via extends."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir) / "base.yaml"
+        child_path = Path(tmpdir) / "child.yaml"
+
+        # Write base config
+        base_path.write_text("""
+executor:
+  adapter: codex_mcp
+  timeout: 3600
+orchestrator:
+  lm: gpt-5-mini
+  max_steps: 30
+""")
+
+        # Write child config that extends base
+        child_path.write_text("""
+extends: base.yaml
+orchestrator:
+  lm: gpt-5
+  prompt: prompts/aggressive.yaml
+""")
+
+        config = load_yaml_config(child_path)
+        assert config["executor"]["adapter"] == "codex_mcp"
+        assert config["executor"]["timeout"] == 3600
+        assert config["orchestrator"]["lm"] == "gpt-5"  # overridden
+        assert config["orchestrator"]["max_steps"] == 30  # inherited
+        assert config["orchestrator"]["prompt"] == "prompts/aggressive.yaml"  # new
+
+
+def test_load_config_full_chain():
+    """Test full config loading with precedence."""
+    import os
+
+    # Save and clear WEAVE_PROJECT to test config precedence
+    orig_weave = os.environ.pop("WEAVE_PROJECT", None)
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir = Path(tmpdir)
+
+            # Write config.toml
+            toml_path = tmpdir / "config.toml"
+            toml_path.write_text("""
+[weave]
+project = "my-weave-project"
+
+[executor]
+adapter = "codex_mcp"
+""")
+
+            # Write yaml config
+            yaml_path = tmpdir / "experiment.yaml"
+            yaml_path.write_text("""
+orchestrator:
+  lm: claude-sonnet
+  max_steps: 100
+""")
+
+            # Load with override (use non-existent env_path to prevent loading cwd's .env)
+            config = load_config(
+                config_path=yaml_path,
+                toml_path=toml_path,
+                env_path=tmpdir / ".env.nonexistent",
+                overrides=["orchestrator.max_steps=200"],
+            )
+
+            # Check precedence: override > yaml > toml > default
+            assert config.weave.project == "my-weave-project"  # from toml
+            assert config.executor.adapter == "codex_mcp"  # from toml
+            assert config.orchestrator.lm == "claude-sonnet"  # from yaml
+            assert config.orchestrator.max_steps == 200  # from override
+    finally:
+        # Restore WEAVE_PROJECT
+        if orig_weave is not None:
+            os.environ["WEAVE_PROJECT"] = orig_weave
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

zwarm/core/test_models.py

@@ -0,0 +1,265 @@
+"""Tests for core models and state management."""
+
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from zwarm.core.models import (
+    ConversationSession,
+    Event,
+    Message,
+    SessionMode,
+    SessionStatus,
+    Task,
+    TaskStatus,
+    event_session_completed,
+    event_session_started,
+    event_task_created,
+)
+from zwarm.core.state import StateManager
+
+
+class TestMessage:
+    def test_create_message(self):
+        msg = Message(role="user", content="Hello")
+        assert msg.role == "user"
+        assert msg.content == "Hello"
+        assert msg.timestamp is not None
+
+    def test_message_serialization(self):
+        msg = Message(role="assistant", content="Hi there")
+        data = msg.to_dict()
+        restored = Message.from_dict(data)
+        assert restored.role == msg.role
+        assert restored.content == msg.content
+
+
+class TestConversationSession:
+    def test_create_session(self):
+        session = ConversationSession(
+            adapter="codex_mcp",
+            mode=SessionMode.SYNC,
+            task_description="Test task",
+        )
+        assert session.id is not None
+        assert session.adapter == "codex_mcp"
+        assert session.mode == SessionMode.SYNC
+        assert session.status == SessionStatus.ACTIVE
+
+    def test_add_message(self):
+        session = ConversationSession()
+        session.add_message("user", "Hello")
+        session.add_message("assistant", "Hi there")
+        assert len(session.messages) == 2
+        assert session.messages[0].role == "user"
+        assert session.messages[1].role == "assistant"
+
+    def test_complete_session(self):
+        session = ConversationSession()
+        session.complete("Done!")
+        assert session.status == SessionStatus.COMPLETED
+        assert session.completed_at is not None
+        assert session.exit_message == "Done!"
+
+    def test_fail_session(self):
+        session = ConversationSession()
+        session.fail("Error occurred")
+        assert session.status == SessionStatus.FAILED
+        assert session.exit_message == "Error occurred"
+
+    def test_session_serialization(self):
+        session = ConversationSession(
+            adapter="claude_code",
+            mode=SessionMode.ASYNC,
+            task_description="Build feature",
+            model="claude-sonnet",
+        )
+        session.add_message("user", "Start")
+        session.conversation_id = "conv-123"
+
+        data = session.to_dict()
+        restored = ConversationSession.from_dict(data)
+
+        assert restored.id == session.id
+        assert restored.adapter == "claude_code"
+        assert restored.mode == SessionMode.ASYNC
+        assert restored.conversation_id == "conv-123"
+        assert len(restored.messages) == 1
+
+
+class TestTask:
+    def test_create_task(self):
+        task = Task(description="Fix the bug")
+        assert task.id is not None
+        assert task.status == TaskStatus.PENDING
+
+    def test_task_lifecycle(self):
+        task = Task(description="Implement feature")
+
+        # Start
+        task.start("session-123")
+        assert task.status == TaskStatus.IN_PROGRESS
+        assert task.session_id == "session-123"
+
+        # Complete
+        task.complete("Feature implemented")
+        assert task.status == TaskStatus.COMPLETED
+        assert task.result == "Feature implemented"
+
+    def test_task_serialization(self):
+        task = Task(description="Test task", parent_task_id="parent-123")
+        task.start("session-456")
+
+        data = task.to_dict()
+        restored = Task.from_dict(data)
+
+        assert restored.id == task.id
+        assert restored.description == "Test task"
+        assert restored.parent_task_id == "parent-123"
+        assert restored.session_id == "session-456"
+
+
+class TestEvent:
+    def test_create_event(self):
+        event = Event(
+            kind="test_event",
+            session_id="session-123",
+            payload={"key": "value"},
+        )
+        assert event.id is not None
+        assert event.kind == "test_event"
+
+    def test_event_factories(self):
+        session = ConversationSession(task_description="Test")
+        event = event_session_started(session)
+        assert event.kind == "session_started"
+        assert event.session_id == session.id
+
+        task = Task(description="Do something")
+        event = event_task_created(task)
+        assert event.kind == "task_created"
+        assert event.task_id == task.id
+
+
+class TestStateManager:
+    def test_init_creates_directories(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            state_dir = Path(tmpdir) / ".zwarm"
+            manager = StateManager(state_dir)
+            manager.init()
+
+            assert state_dir.exists()
+            assert (state_dir / "sessions").exists()
+            assert (state_dir / "orchestrator").exists()
+            assert (state_dir / "events.jsonl").exists()
+
+    def test_session_crud(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(Path(tmpdir) / ".zwarm")
+            manager.init()
+
+            # Add
+            session = ConversationSession(task_description="Test")
+            manager.add_session(session)
+
+            # Get
+            retrieved = manager.get_session(session.id)
+            assert retrieved is not None
+            assert retrieved.task_description == "Test"
+
+            # Update
+            session.add_message("user", "Hello")
+            manager.update_session(session)
+
+            # List
+            sessions = manager.list_sessions()
+            assert len(sessions) == 1
+
+            # Filter by status
+            active = manager.list_sessions(status="active")
+            assert len(active) == 1
+
+    def test_task_crud(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(Path(tmpdir) / ".zwarm")
+            manager.init()
+
+            # Add
+            task = Task(description="Build feature")
+            manager.add_task(task)
+
+            # Get
+            retrieved = manager.get_task(task.id)
+            assert retrieved is not None
+
+            # Update
+            task.start("session-123")
+            manager.update_task(task)
+
+            # List
+            tasks = manager.list_tasks(status="in_progress")
+            assert len(tasks) == 1
+
+    def test_event_logging(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(Path(tmpdir) / ".zwarm")
+            manager.init()
+
+            # Log events
+            session = ConversationSession()
+            manager.log_event(event_session_started(session))
+            manager.log_event(event_session_completed(session))
+
+            # Read events
+            events = manager.get_events()
+            assert len(events) == 2
+            assert events[0].kind == "session_completed"  # Most recent first
+
+            # Filter by session
+            events = manager.get_events(session_id=session.id)
+            assert len(events) == 2
+
+            # Filter by kind
+            events = manager.get_events(kind="session_started")
+            assert len(events) == 1
+
+    def test_orchestrator_messages(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            manager = StateManager(Path(tmpdir) / ".zwarm")
+            manager.init()
+
+            messages = [
+                {"role": "system", "content": "You are an orchestrator"},
+                {"role": "user", "content": "Build a feature"},
+            ]
+            manager.save_orchestrator_messages(messages)
+
+            loaded = manager.load_orchestrator_messages()
+            assert len(loaded) == 2
+            assert loaded[0]["role"] == "system"
+
+    def test_state_persistence(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            state_dir = Path(tmpdir) / ".zwarm"
+
+            # Create and save state
+            manager1 = StateManager(state_dir)
+            manager1.init()
+
+            session = ConversationSession(task_description="Persistent session")
+            manager1.add_session(session)
+
+            task = Task(description="Persistent task")
+            manager1.add_task(task)
+
+            # Load in new manager
+            manager2 = StateManager(state_dir)
+            manager2.load()
+
+            assert manager2.get_session(session.id) is not None
+            assert manager2.get_task(task.id) is not None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])