zwarm 0.1.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

zwarm/core/test_compact.py

@@ -0,0 +1,266 @@
+"""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

zwarm/orchestrator.py

@@ -24,12 +24,18 @@ from wbal.helper import TOOL_CALL_TYPE, format_openai_tool_response
 from zwarm.adapters.base import ExecutorAdapter
 from zwarm.adapters.claude_code import ClaudeCodeAdapter
 from zwarm.adapters.codex_mcp import CodexMCPAdapter
+from zwarm.core.compact import compact_messages, should_compact
 from zwarm.core.config import ZwarmConfig, load_config
 from zwarm.core.environment import OrchestratorEnv
 from zwarm.core.models import ConversationSession
 from zwarm.core.state import StateManager
 from zwarm.prompts import get_orchestrator_prompt
-from zwarm.watchers import WatcherManager, WatcherContext, WatcherAction, build_watcher_manager
+from zwarm.watchers import (
+    WatcherManager,
+    WatcherContext,
+    WatcherAction,
+    build_watcher_manager,
+)
 
 
 class Orchestrator(YamlAgent):
@@ -49,10 +55,12 @@ class Orchestrator(YamlAgent):
     working_dir: Path = Field(default_factory=Path.cwd)
 
     # Load tools from modules (delegation + bash for verification)
-    agent_tool_modules: list[str] = Field(default=[
-        "zwarm.tools.delegation",
-        "wbal.tools.bash",
-    ])
+    agent_tool_modules: list[str] = Field(
+        default=[
+            "zwarm.tools.delegation",
+            "wbal.tools.bash",
+        ]
+    )
 
     # State management
     _state: StateManager = PrivateAttr()
@@ -60,6 +68,12 @@ class Orchestrator(YamlAgent):
     _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
+    _executor_usage: dict[str, int] = PrivateAttr(default_factory=lambda: {
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "total_tokens": 0,
+    })
 
     def model_post_init(self, __context: Any) -> None:
         """Initialize state and adapters after model creation."""
@@ -80,17 +94,29 @@ class Orchestrator(YamlAgent):
 
         # Initialize watchers if configured
         if self.config.watchers.enabled:
-            self._watcher_manager = build_watcher_manager({
-                "watchers": [
-                    {"name": w.name, "enabled": w.enabled, "config": w.config}
-                    for w in self.config.watchers.watchers
-                ]
-            })
+            self._watcher_manager = build_watcher_manager(
+                {
+                    "watchers": [
+                        {"name": w.name, "enabled": w.enabled, "config": w.config}
+                        for w in self.config.watchers.watchers
+                    ]
+                }
+            )
 
         # Link sessions to environment for observe()
-        if hasattr(self.env, 'set_sessions'):
+        if hasattr(self.env, "set_sessions"):
             self.env.set_sessions(self._sessions)
 
+        # Set budget limits in environment
+        if hasattr(self.env, "set_budget"):
+            # Extract budget from watcher config if available
+            max_sessions = None
+            for w in self.config.watchers.watchers:
+                if w.name == "budget" and w.config:
+                    max_sessions = w.config.get("max_sessions")
+                    break
+            self.env.set_budget(max_sessions=max_sessions)
+
     @property
     def state(self) -> StateManager:
         """Access state manager."""
@@ -99,22 +125,97 @@ class Orchestrator(YamlAgent):
     def _get_adapter(self, name: str) -> ExecutorAdapter:
         """Get or create an adapter by name."""
         if name not in self._adapters:
+            # Get model from config (adapters have their own defaults if None)
+            model = self.config.executor.model
             if name == "codex_mcp":
-                self._adapters[name] = CodexMCPAdapter()
+                self._adapters[name] = CodexMCPAdapter(model=model)
             elif name == "claude_code":
-                self._adapters[name] = ClaudeCodeAdapter()
+                self._adapters[name] = ClaudeCodeAdapter(model=model)
             else:
                 raise ValueError(f"Unknown adapter: {name}")
         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()
+
     def save_state(self) -> None:
         """Save orchestrator state for resume."""
         self._state.save_orchestrator_messages(self.messages)
 
     def load_state(self) -> None:
-        """Load orchestrator state for resume."""
-        self.messages = self._state.load_orchestrator_messages()
-        self._resumed = True
+        """Load orchestrator state for resume.
+
+        Only marks as resumed if we actually loaded non-empty messages.
+        This prevents the resume message from being injected before the
+        system prompt when there's no saved state to resume from.
+        """
+        loaded_messages = self._state.load_orchestrator_messages()
+        if loaded_messages:
+            self.messages = loaded_messages
+            self._resumed = True
+        # If no messages were saved, don't set _resumed - start fresh
+
+    def _maybe_compact(self) -> bool:
+        """
+        Check if compaction is needed and compact if so.
+
+        Returns True if compaction was performed.
+        """
+        compact_config = self.config.orchestrator.compaction
+        if not compact_config.enabled:
+            return False
+
+        # Check if we should compact
+        if not should_compact(
+            self.messages,
+            max_tokens=compact_config.max_tokens,
+            threshold_pct=compact_config.threshold_pct,
+        ):
+            return False
+
+        # Perform compaction
+        result = compact_messages(
+            self.messages,
+            keep_first_n=compact_config.keep_first_n,
+            keep_last_n=compact_config.keep_last_n,
+            max_tokens=compact_config.max_tokens,
+            target_token_pct=compact_config.target_pct,
+        )
+
+        if result.was_compacted:
+            self.messages = result.messages
+
+            # Log compaction event
+            from zwarm.core.models import Event
+            self._state.log_event(Event(
+                kind="context_compacted",
+                payload={
+                    "step": self._step_count,
+                    "original_count": result.original_count,
+                    "new_count": len(result.messages),
+                    "removed_count": result.removed_count,
+                },
+            ))
+
+            return True
+
+        return False
 
     def _inject_resume_message(self) -> None:
         """Inject a system message about resumed state."""
@@ -124,7 +225,9 @@ class Orchestrator(YamlAgent):
         # Build list of old sessions
         old_sessions = []
         for sid, session in self._sessions.items():
-            old_sessions.append(f"  - {sid[:8]}... ({session.adapter}, {session.status.value})")
+            old_sessions.append(
+                f"  - {sid[:8]}... ({session.adapter}, {session.status.value})"
+            )
 
         session_info = "\n".join(old_sessions) if old_sessions else "  (none)"
 
@@ -139,31 +242,67 @@ Previous sessions (now stale):
 
 You must start NEW sessions with delegate() if you need to continue work. Do NOT try to use converse() or check_session() with the old session IDs - they will fail.
 
-Continue with your task from where you left off."""
+Continue with your task from where you left off.""",
         }
 
         self.messages.append(resume_msg)
         self._resumed = False  # Only inject once
 
+    def perceive(self) -> None:
+        """
+        Override perceive to refresh environment observation each step.
+
+        The base YamlAgent only adds env.observe() on step 0. We need to
+        update it each step to show current progress, sessions, etc.
+        """
+        # Let base class do initial setup
+        super().perceive()
+
+        # Update environment observation
+        env_obs = (self.env.observe() or "").strip()
+        if not env_obs:
+            return
+
+        # Find and update existing env observation, or append new one
+        # Look for a system message containing our markers
+        env_marker = "## Progress"  # Our env observation has this
+
+        for i, msg in enumerate(self.messages):
+            if msg.get("role") == "system" and env_marker in msg.get("content", ""):
+                # Update in place
+                self.messages[i]["content"] = env_obs
+                return
+
+        # Not found - append as new system message (shouldn't happen after step 0)
+        self.messages.append({"role": "system", "content": env_obs})
+
+    @weave.op()
     def _run_watchers(self) -> WatcherAction:
         """Run watchers and return the action to take."""
         if not self._watcher_manager:
             return WatcherAction.CONTINUE
 
         # Build watcher context
+        task = getattr(self.env, "task", "") if self.env else ""
+        events = [e.to_dict() for e in self.state.get_events(limit=200)]
         ctx = WatcherContext(
+            task=task,
             step=self._step_count,
+            max_steps=self.maxSteps,
             messages=self.messages,
-            sessions={sid: s.to_dict() for sid, s in self._sessions.items()},
-            task=self.env.task if hasattr(self.env, 'task') else "",
+            sessions=[s.to_dict() for s in self._sessions.values()],
+            events=events,
+            working_dir=str(self.working_dir.absolute()) if self.working_dir else None,
             metadata={
-                "max_steps": self.maxSteps,
-                "config": self.config.to_dict() if hasattr(self.config, 'to_dict') else {},
+                "config": self.config.to_dict()
+                if hasattr(self.config, "to_dict")
+                else {},
             },
         )
 
         # Run watchers synchronously (they're async internally)
         import asyncio
+
         try:
             loop = asyncio.get_running_loop()
         except RuntimeError:
@@ -172,18 +311,37 @@ Continue with your task from where you left off."""
         if loop and loop.is_running():
             # We're in an async context, create a task
             import concurrent.futures
+
             with concurrent.futures.ThreadPoolExecutor() as pool:
-                result = pool.submit(asyncio.run, self._watcher_manager.observe(ctx)).result()
+                result = pool.submit(
+                    asyncio.run, self._watcher_manager.observe(ctx)
+                ).result()
         else:
             result = asyncio.run(self._watcher_manager.observe(ctx))
 
+        # Log watcher execution to events
+        from zwarm.core.models import Event
+        watcher_names = [w.name for w in self.config.watchers.watchers if w.enabled]
+        self.state.log_event(Event(
+            kind="watchers_run",
+            payload={
+                "step": self._step_count,
+                "watchers": watcher_names,
+                "action": result.action.value,
+                "triggered_by": result.metadata.get("triggered_by"),
+                "reason": result.metadata.get("reason"),
+            },
+        ))
+
         # Handle watcher result
         if result.action == WatcherAction.NUDGE and result.guidance:
             # Inject guidance as a system message
-            self.messages.append({
-                "role": "user",
-                "content": f"[WATCHER: {result.metadata.get('triggered_by', 'unknown')}] {result.guidance}"
-            })
+            self.messages.append(
+                {
+                    "role": "user",
+                    "content": f"[WATCHER: {result.metadata.get('triggered_by', 'unknown')}] {result.guidance}",
+                }
+            )
 
         return result.action
 
@@ -200,20 +358,19 @@ Continue with your task from where you left off."""
         if self._last_response is None:
             return []
 
-        output = getattr(self._last_response, 'output', None)
+        output = getattr(self._last_response, "output", None)
         if output is None:
             return []
 
         # Extract tool calls
         tool_calls = [
-            item for item in output
-            if getattr(item, 'type', None) == TOOL_CALL_TYPE
+            item for item in output if getattr(item, "type", None) == TOOL_CALL_TYPE
         ]
 
         # If no tool calls, handle text output
         if not tool_calls:
-            output_text = getattr(self._last_response, 'output_text', '')
-            if output_text and hasattr(self.env, 'output_handler'):
+            output_text = getattr(self._last_response, "output_text", "")
+            if output_text and hasattr(self.env, "output_handler"):
                 self.env.output_handler(output_text)
             return []
 
@@ -221,9 +378,9 @@ Continue with your task from where you left off."""
         tool_results: list[tuple[dict[str, Any], Any]] = []
 
         for tc in tool_calls:
-            tc_name = getattr(tc, 'name', '')
-            tc_args_raw = getattr(tc, 'arguments', '{}')
-            tc_id = getattr(tc, 'call_id', '')
+            tc_name = getattr(tc, "name", "")
+            tc_args_raw = getattr(tc, "arguments", "{}")
+            tc_id = getattr(tc, "call_id", "")
 
             # Parse arguments
             if isinstance(tc_args_raw, str):
@@ -271,14 +428,37 @@ Continue with your task from where you left off."""
             - tool_call_info: {"name": str, "args": dict, "call_id": str}
             - result: The tool output (any type)
         """
+        # Check for context compaction before perceive
+        # This prevents context overflow on long-running tasks
+        self._maybe_compact()
+
+        # Update environment with current progress before perceive
+        if hasattr(self.env, "update_progress"):
+            executor_usage = self.get_executor_usage()
+            self.env.update_progress(
+                step_count=self._step_count,
+                max_steps=self.maxSteps,
+                total_tokens=self._total_tokens,
+                executor_tokens=executor_usage.get("total_tokens", 0),
+            )
+
         self.perceive()
         self.invoke()
+
+        # Track cumulative token usage from the API response
+        if self._last_response and hasattr(self._last_response, "usage"):
+            usage = self._last_response.usage
+            if usage:
+                self._total_tokens += getattr(usage, "total_tokens", 0)
+
         tool_results = self.do()
         self._step_count += 1
         return tool_results
 
     @weave.op()
-    def run(self, task: str | None = None, max_steps: int | None = None) -> dict[str, Any]:
+    def run(
+        self, task: str | None = None, max_steps: int | None = None
+    ) -> dict[str, Any]:
         """
         Run the orchestrator until stop condition is met.
 
@@ -299,8 +479,9 @@ Continue with your task from where you left off."""
         if max_steps is not None:
             self.maxSteps = max_steps
 
-        # Reset step counter
+        # Reset counters
         self._step_count = 0
+        self._total_tokens = 0
 
         # Inject resume message if we were resumed
         self._inject_resume_message()
@@ -402,4 +583,6 @@ def build_orchestrator(
 
 def _build_system_prompt(config: ZwarmConfig, working_dir: Path | None = None) -> str:
     """Build the orchestrator system prompt."""
-    return get_orchestrator_prompt(working_dir=str(working_dir) if working_dir else None)
+    return get_orchestrator_prompt(
+        working_dir=str(working_dir) if working_dir else None
+    )