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

zwarm/adapters/claude_code.py

@@ -33,10 +33,46 @@ class ClaudeCodeAdapter(ExecutorAdapter):
     """
 
     name = "claude_code"
+    DEFAULT_MODEL = "claude-sonnet-4-5-20250514"  # Best balance of speed and capability
 
     def __init__(self, model: str | None = None):
-        self._model = model
+        self._model = model or self.DEFAULT_MODEL
         self._sessions: dict[str, str] = {}  # session_id -> claude session_id
+        # Cumulative token usage for cost tracking
+        self._total_usage: dict[str, int] = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "cache_creation_input_tokens": 0,
+            "cache_read_input_tokens": 0,
+            "total_tokens": 0,
+        }
+
+    def _accumulate_usage(self, usage: dict[str, Any]) -> None:
+        """Add usage to cumulative totals."""
+        if not usage:
+            return
+        for key in self._total_usage:
+            self._total_usage[key] += usage.get(key, 0)
+
+    def _extract_usage(self, output: dict) -> dict[str, int]:
+        """Extract token usage from claude CLI JSON output."""
+        usage = {}
+        # Claude CLI may include usage in various formats
+        if "usage" in output:
+            usage = output["usage"]
+        elif "cost_usd" in output:
+            # Alternative: estimate from cost if available
+            usage["cost_usd"] = output["cost_usd"]
+        # Also check for token counts in the output
+        for key in ["input_tokens", "output_tokens", "total_tokens"]:
+            if key in output:
+                usage[key] = output[key]
+        return usage
+
+    @property
+    def total_usage(self) -> dict[str, int]:
+        """Get cumulative token usage across all calls."""
+        return self._total_usage.copy()
 
     @weave.op()
     async def _call_claude(
@@ -75,11 +111,14 @@ class ClaudeCodeAdapter(ExecutorAdapter):
 
         response_text = self._extract_response(result.stdout, result.stderr)
 
-        # Try to get session ID from JSON output
+        # Try to get session ID and usage from JSON output
         session_id = None
+        usage = {}
         try:
             output = json.loads(result.stdout)
             session_id = output.get("session_id")
+            usage = self._extract_usage(output)
+            self._accumulate_usage(usage)
         except (json.JSONDecodeError, TypeError):
             pass
 
@@ -87,6 +126,8 @@ class ClaudeCodeAdapter(ExecutorAdapter):
             "response": response_text,
             "session_id": session_id,
             "exit_code": result.returncode,
+            "usage": usage,
+            "total_usage": self.total_usage,
         }
 
     @weave.op()
@@ -126,16 +167,21 @@ class ClaudeCodeAdapter(ExecutorAdapter):
 
         response_text = self._extract_response(result.stdout, result.stderr)
 
-        # Try to get session ID from JSON output
+        # Try to get session ID and usage from JSON output
         new_session_id = None
+        usage = {}
         try:
             output = json.loads(result.stdout)
             new_session_id = output.get("session_id")
+            usage = self._extract_usage(output)
+            self._accumulate_usage(usage)
         except (json.JSONDecodeError, TypeError):
             pass
 
         return {
             "response": response_text,
+            "usage": usage,
+            "total_usage": self.total_usage,
             "session_id": new_session_id or session_id,
             "exit_code": result.returncode,
         }
@@ -175,6 +221,9 @@ class ClaudeCodeAdapter(ExecutorAdapter):
             session.add_message("user", task)
             session.add_message("assistant", result["response"])
 
+            # Track token usage on the session
+            session.add_usage(result.get("usage", {}))
+
         else:
             # Async mode: run in background
             cmd = ["claude", "-p", "--output-format", "json"]
@@ -223,6 +272,9 @@ class ClaudeCodeAdapter(ExecutorAdapter):
         session.add_message("user", message)
         session.add_message("assistant", response_text)
 
+        # Track token usage on the session
+        session.add_usage(result.get("usage", {}))
+
         return response_text
 
     async def check_status(