zwarm 2.3.5__py3-none-any.whl

zwarm/core/state.py

@@ -0,0 +1,355 @@
+"""
+Flat-file state management for zwarm.
+
+State structure (with instance isolation):
+.zwarm/
+├── instances.json              # Registry of all instances
+└── instances/
+    └── <instance-id>/
+        ├── state.json          # Current state (sessions, tasks)
+        ├── events.jsonl        # Append-only event log
+        ├── sessions/
+        │   └── <session-id>/
+        │       ├── messages.json
+        │       └── output.log
+        └── orchestrator/
+            └── messages.json   # Orchestrator's message history (for resume)
+
+Legacy structure (single instance, for backwards compat):
+.zwarm/
+├── state.json
+├── events.jsonl
+└── ...
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from .models import ConversationSession, Event, Task
+
+
+# --- Instance Registry ---
+
+def get_instances_registry_path(base_dir: Path | str = ".zwarm") -> Path:
+    """Get path to the instances registry file."""
+    return Path(base_dir) / "instances.json"
+
+
+def list_instances(base_dir: Path | str = ".zwarm") -> list[dict[str, Any]]:
+    """List all registered instances."""
+    registry_path = get_instances_registry_path(base_dir)
+    if not registry_path.exists():
+        return []
+    try:
+        return json.loads(registry_path.read_text()).get("instances", [])
+    except (json.JSONDecodeError, KeyError):
+        return []
+
+
+def register_instance(
+    instance_id: str,
+    name: str | None = None,
+    task: str | None = None,
+    base_dir: Path | str = ".zwarm",
+) -> None:
+    """Register an instance in the global registry."""
+    base = Path(base_dir)
+    base.mkdir(parents=True, exist_ok=True)
+
+    registry_path = get_instances_registry_path(base_dir)
+
+    # Load existing registry
+    if registry_path.exists():
+        try:
+            registry = json.loads(registry_path.read_text())
+        except json.JSONDecodeError:
+            registry = {"instances": []}
+    else:
+        registry = {"instances": []}
+
+    # Check if instance already registered
+    existing_ids = {inst["id"] for inst in registry["instances"]}
+    if instance_id in existing_ids:
+        # Update existing entry
+        for inst in registry["instances"]:
+            if inst["id"] == instance_id:
+                inst["updated_at"] = datetime.now().isoformat()
+                inst["status"] = "active"
+                if name:
+                    inst["name"] = name
+                if task:
+                    inst["task"] = task[:100]  # Truncate
+                break
+    else:
+        # Add new entry
+        registry["instances"].append({
+            "id": instance_id,
+            "name": name or instance_id[:8],
+            "task": (task[:100] if task else None),
+            "created_at": datetime.now().isoformat(),
+            "updated_at": datetime.now().isoformat(),
+            "status": "active",
+        })
+
+    registry_path.write_text(json.dumps(registry, indent=2))
+
+
+def update_instance_status(
+    instance_id: str,
+    status: str,
+    base_dir: Path | str = ".zwarm",
+) -> None:
+    """Update an instance's status in the registry."""
+    registry_path = get_instances_registry_path(base_dir)
+    if not registry_path.exists():
+        return
+
+    try:
+        registry = json.loads(registry_path.read_text())
+    except json.JSONDecodeError:
+        return
+
+    for inst in registry.get("instances", []):
+        if inst["id"] == instance_id:
+            inst["status"] = status
+            inst["updated_at"] = datetime.now().isoformat()
+            break
+
+    registry_path.write_text(json.dumps(registry, indent=2))
+
+
+def get_instance_state_dir(
+    instance_id: str | None = None,
+    base_dir: Path | str = ".zwarm",
+) -> Path:
+    """
+    Get the state directory for an instance.
+
+    If instance_id is None, returns the legacy path for backwards compat.
+    """
+    base = Path(base_dir)
+    if instance_id is None:
+        return base  # Legacy: .zwarm/
+    return base / "instances" / instance_id
+
+
+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.
+    With instance isolation: .zwarm/instances/<instance-id>/
+    Legacy (no instance): .zwarm/
+
+    This enables:
+    - Git-backed history
+    - Easy debugging (just read the files)
+    - Resume from previous state
+    - Multiple concurrent orchestrators (with instance isolation)
+    """
+
+    def __init__(
+        self,
+        state_dir: Path | str = ".zwarm",
+        instance_id: str | None = None,
+    ):
+        self.base_dir = Path(state_dir)
+        self.instance_id = instance_id
+
+        # Resolve actual state directory
+        if instance_id:
+            self.state_dir = get_instance_state_dir(instance_id, self.base_dir)
+        else:
+            self.state_dir = self.base_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_compact.py

@@ -0,0 +1,312 @@
+"""Tests for the compact module."""
+
+import pytest
+
+from zwarm.core.compact import (
+    compact_messages,
+    estimate_tokens,
+    find_tool_groups,
+    should_compact,
+)
+
+
+class TestEstimateTokens:
+    def test_simple_messages(self):
+        """Estimate tokens for simple text messages."""
+        messages = [
+            {"role": "user", "content": "Hello world"},  # 11 chars
+            {"role": "assistant", "content": "Hi there!"},  # 9 chars
+        ]
+        # ~20 chars / 4 = ~5 tokens
+        tokens = estimate_tokens(messages)
+        assert tokens == 5
+
+    def test_empty_messages(self):
+        """Empty messages return 0 tokens."""
+        assert estimate_tokens([]) == 0
+
+    def test_messages_with_tool_calls(self):
+        """Tool calls add to token count."""
+        messages = [
+            {
+                "role": "assistant",
+                "content": "Let me check",
+                "tool_calls": [
+                    {"function": {"name": "read", "arguments": '{"path": "/foo/bar"}'}}
+                ],
+            }
+        ]
+        tokens = estimate_tokens(messages)
+        assert tokens > 0
+
+
+class TestFindToolGroups:
+    def test_no_tool_calls(self):
+        """No tool groups in simple conversation."""
+        messages = [
+            {"role": "system", "content": "You are helpful"},
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi!"},
+        ]
+        groups = find_tool_groups(messages)
+        assert groups == []
+
+    def test_openai_format_tool_call(self):
+        """Detect OpenAI-style tool call groups."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Read file"},
+            {
+                "role": "assistant",
+                "content": "Reading...",
+                "tool_calls": [{"id": "tc1", "function": {"name": "read"}}],
+            },
+            {"role": "tool", "tool_call_id": "tc1", "content": "file contents"},
+            {"role": "assistant", "content": "Here's the file"},
+        ]
+        groups = find_tool_groups(messages)
+        assert groups == [(2, 3)]  # Assistant with tool_calls + tool response
+
+    def test_multiple_tool_responses(self):
+        """Group includes all consecutive tool responses."""
+        messages = [
+            {"role": "user", "content": "Do things"},
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {"id": "tc1", "function": {"name": "a"}},
+                    {"id": "tc2", "function": {"name": "b"}},
+                ],
+            },
+            {"role": "tool", "tool_call_id": "tc1", "content": "result1"},
+            {"role": "tool", "tool_call_id": "tc2", "content": "result2"},
+            {"role": "assistant", "content": "Done"},
+        ]
+        groups = find_tool_groups(messages)
+        assert groups == [(1, 3)]  # Indices 1, 2, 3 form one group
+
+    def test_anthropic_format_tool_use(self):
+        """Detect Anthropic-style tool_use content blocks."""
+        messages = [
+            {"role": "user", "content": "Read file"},
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "text", "text": "Reading..."},
+                    {"type": "tool_use", "id": "tu1", "name": "read", "input": {}},
+                ],
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "tool_result", "tool_use_id": "tu1", "content": "data"},
+                ],
+            },
+            {"role": "assistant", "content": "Got it"},
+        ]
+        groups = find_tool_groups(messages)
+        assert groups == [(1, 2)]  # Assistant with tool_use + user with tool_result
+
+
+class TestCompactMessages:
+    def test_no_compaction_needed_few_messages(self):
+        """Don't compact if we have fewer messages than keep thresholds."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Task"},
+            {"role": "assistant", "content": "Response"},
+        ]
+        result = compact_messages(messages, keep_first_n=2, keep_last_n=2)
+        assert not result.was_compacted
+        assert result.messages == messages
+        assert "Too few" in result.preserved_reason
+
+    def test_compacts_middle_messages(self):
+        """Remove messages from the middle, keeping first and last."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Task"},
+            {"role": "assistant", "content": "Step 1"},
+            {"role": "user", "content": "Continue"},
+            {"role": "assistant", "content": "Step 2"},
+            {"role": "user", "content": "More"},
+            {"role": "assistant", "content": "Step 3"},
+            {"role": "user", "content": "Final"},
+            {"role": "assistant", "content": "Done"},
+        ]
+        result = compact_messages(messages, keep_first_n=2, keep_last_n=2)
+
+        assert result.was_compacted
+        assert result.removed_count > 0
+        # First 2 and last 2 should be preserved
+        assert result.messages[0]["content"] == "System"
+        assert result.messages[1]["content"] == "Task"
+        assert result.messages[-1]["content"] == "Done"
+        assert result.messages[-2]["content"] == "Final"
+
+    def test_preserves_tool_call_pairs(self):
+        """Never split tool call from its response."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Task"},
+            {"role": "assistant", "content": "Old message 1"},
+            {"role": "assistant", "content": "Old message 2"},
+            {
+                "role": "assistant",
+                "content": "Calling tool",
+                "tool_calls": [{"id": "tc1", "function": {"name": "test"}}],
+            },
+            {"role": "tool", "tool_call_id": "tc1", "content": "Tool result"},
+            {"role": "assistant", "content": "Recent 1"},
+            {"role": "user", "content": "Recent 2"},
+        ]
+        result = compact_messages(messages, keep_first_n=2, keep_last_n=2)
+
+        # The tool call pair should either both be kept or both removed
+        has_tool_call = any(m.get("tool_calls") for m in result.messages)
+        has_tool_response = any(m.get("role") == "tool" for m in result.messages)
+
+        # They should match - either both present or both absent
+        assert has_tool_call == has_tool_response
+
+    def test_adds_compaction_marker(self):
+        """Add a marker message when compaction occurs."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Task"},
+        ] + [{"role": "assistant", "content": f"Msg {i}"} for i in range(20)]
+
+        result = compact_messages(messages, keep_first_n=2, keep_last_n=3)
+
+        if result.was_compacted:
+            # Should have a system message about compaction
+            marker_msgs = [
+                m for m in result.messages
+                if m.get("role") == "system" and "compacted" in m.get("content", "").lower()
+            ]
+            assert len(marker_msgs) == 1
+
+    def test_token_based_compaction(self):
+        """Compact based on token threshold."""
+        # Create messages that exceed token limit
+        messages = [
+            {"role": "system", "content": "System prompt " * 100},
+            {"role": "user", "content": "Task " * 100},
+        ] + [
+            {"role": "assistant", "content": f"Response {i} " * 50}
+            for i in range(10)
+        ]
+
+        # Should not compact if under limit
+        result_under = compact_messages(messages, max_tokens=100000)
+        # Might or might not compact depending on estimate
+
+        # Should compact if over limit
+        result_over = compact_messages(messages, max_tokens=100, target_token_pct=0.5)
+        # With such a low limit, should definitely try to compact
+        assert result_over.original_count == len(messages)
+
+
+class TestShouldCompact:
+    def test_under_threshold(self):
+        """Don't compact when under threshold."""
+        messages = [{"role": "user", "content": "Hello"}]
+        assert not should_compact(messages, max_tokens=1000, threshold_pct=0.85)
+
+    def test_over_threshold(self):
+        """Compact when over threshold."""
+        messages = [{"role": "user", "content": "x" * 4000}]  # ~1000 tokens
+        assert should_compact(messages, max_tokens=500, threshold_pct=0.85)
+
+
+class TestEdgeCases:
+    def test_all_tool_calls(self):
+        """Handle conversation that's mostly tool calls."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Task"},
+        ]
+        # Add many tool call pairs
+        for i in range(5):
+            messages.append({
+                "role": "assistant",
+                "tool_calls": [{"id": f"tc{i}", "function": {"name": "test"}}],
+            })
+            messages.append({"role": "tool", "tool_call_id": f"tc{i}", "content": f"result{i}"})
+
+        messages.append({"role": "assistant", "content": "Final"})
+
+        result = compact_messages(messages, keep_first_n=2, keep_last_n=1)
+
+        # Should still produce valid output
+        assert len(result.messages) > 0
+
+        # Check no orphaned tool calls
+        for i, msg in enumerate(result.messages):
+            if msg.get("tool_calls"):
+                # Next message should be a tool response
+                if i + 1 < len(result.messages):
+                    # Either next is tool response, or this is at the end
+                    pass  # Structural validity checked by not raising
+
+    def test_empty_messages(self):
+        """Handle empty message list."""
+        result = compact_messages([])
+        assert result.messages == []
+        assert not result.was_compacted
+
+    def test_only_system_and_user(self):
+        """Handle minimal conversation."""
+        messages = [
+            {"role": "system", "content": "System"},
+            {"role": "user", "content": "Hello"},
+        ]
+        result = compact_messages(messages, keep_first_n=2, keep_last_n=2)
+        assert not result.was_compacted
+        assert result.messages == messages
+
+
+class TestPydanticModelMessages:
+    """Test handling of Pydantic model messages (not just dicts)."""
+
+    def test_estimate_tokens_with_objects(self):
+        """estimate_tokens should handle objects with attributes."""
+        class MockMessage:
+            def __init__(self, role, content):
+                self.role = role
+                self.content = content
+
+        messages = [
+            MockMessage("user", "Hello world"),
+            MockMessage("assistant", "Hi there!"),
+        ]
+        tokens = estimate_tokens(messages)
+        assert tokens > 0
+
+    def test_should_compact_with_objects(self):
+        """should_compact should handle objects with attributes."""
+        class MockMessage:
+            def __init__(self, role, content):
+                self.role = role
+                self.content = content
+
+        messages = [MockMessage("user", "x" * 4000)]
+        # Should not crash
+        result = should_compact(messages, max_tokens=500, threshold_pct=0.85)
+        assert result is True
+
+    def test_find_tool_groups_with_objects(self):
+        """find_tool_groups should handle objects with attributes."""
+        class MockMessage:
+            def __init__(self, role, content=None, tool_calls=None):
+                self.role = role
+                self.content = content
+                self.tool_calls = tool_calls
+
+        messages = [
+            MockMessage("user", "Task"),
+            MockMessage("assistant", "Done"),
+        ]
+        # Should not crash
+        groups = find_tool_groups(messages)
+        assert groups == []