agentpool 2.1.9__py3-none-any.whl → 2.5.0__py3-none-any.whl

agentpool/agents/claude_code_agent/usage.py

@@ -0,0 +1,242 @@
+"""Claude Code usage limits helper.
+
+Fetches usage information from Anthropic's OAuth API using stored credentials.
+Works on both Linux and macOS.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from pathlib import Path
+import subprocess
+import sys
+
+import anyenv
+import httpx
+from pydantic import BaseModel
+
+
+class UsageLimit(BaseModel):
+    """A single usage limit with utilization percentage and reset time."""
+
+    utilization: float
+    """Utilization percentage (0-100)."""
+
+    resets_at: datetime | None = None
+    """When this limit resets, or None if not set."""
+
+
+class ExtraUsage(BaseModel):
+    """Extra usage information for paid plans."""
+
+    is_enabled: bool = False
+    monthly_limit: float | None = None
+    used_credits: float | None = None
+    utilization: float | None = None
+
+
+class ClaudeCodeUsage(BaseModel):
+    """Claude Code usage limits."""
+
+    five_hour: UsageLimit | None = None
+    """5-hour rolling usage limit."""
+
+    seven_day: UsageLimit | None = None
+    """7-day rolling usage limit."""
+
+    seven_day_opus: UsageLimit | None = None
+    """7-day Opus-specific limit."""
+
+    seven_day_sonnet: UsageLimit | None = None
+    """7-day Sonnet-specific limit."""
+
+    seven_day_oauth_apps: UsageLimit | None = None
+    """7-day OAuth apps limit."""
+
+    extra_usage: ExtraUsage | None = None
+    """Extra usage info for paid plans."""
+
+    def format_table(self) -> str:
+        """Format usage as a readable table."""
+        lines = ["Claude Code Usage Limits", "=" * 50]
+
+        def format_limit(name: str, limit: UsageLimit | None) -> str | None:
+            if limit is None:
+                return None
+            reset_str = ""
+            if limit.resets_at:
+                reset_str = f" (resets {limit.resets_at.strftime('%Y-%m-%d %H:%M UTC')})"
+            return f"{name}: {limit.utilization:.0f}%{reset_str}"
+
+        for name, limit in [
+            ("5-hour", self.five_hour),
+            ("7-day", self.seven_day),
+            ("7-day Opus", self.seven_day_opus),
+            ("7-day Sonnet", self.seven_day_sonnet),
+            ("7-day OAuth Apps", self.seven_day_oauth_apps),
+        ]:
+            formatted = format_limit(name, limit)
+            if formatted:
+                lines.append(formatted)
+
+        if self.extra_usage and self.extra_usage.is_enabled:
+            lines.append("")
+            lines.append("Extra Usage (paid):")
+            if self.extra_usage.utilization is not None:
+                lines.append(f"  Utilization: {self.extra_usage.utilization:.0f}%")
+            if self.extra_usage.used_credits is not None:
+                lines.append(f"  Used credits: {self.extra_usage.used_credits}")
+            if self.extra_usage.monthly_limit is not None:
+                lines.append(f"  Monthly limit: {self.extra_usage.monthly_limit}")
+
+        return "\n".join(lines)
+
+
+def _get_credentials_path() -> Path | None:
+    """Get the path to Claude Code credentials file.
+
+    Returns:
+        Path to credentials file, or None if not found.
+    """
+    # Linux: ~/.claude/.credentials.json
+    linux_path = Path.home() / ".claude" / ".credentials.json"
+    if linux_path.exists():
+        return linux_path
+
+    # macOS: Also check ~/.claude first (newer versions)
+    if linux_path.exists():
+        return linux_path
+
+    return None
+
+
+def _get_access_token_from_file(path: Path) -> str | None:
+    """Read access token from credentials file."""
+    try:
+        data = anyenv.load_json(path.read_text(), return_type=dict)
+        val = data.get("claudeAiOauth", {}).get("accessToken")
+    except (anyenv.JsonLoadError, OSError):
+        return None
+    else:
+        assert isinstance(val, str)
+        return val
+
+
+def _get_access_token_from_keychain() -> str | None:
+    """Read access token from macOS Keychain."""
+    if sys.platform != "darwin":
+        return None
+
+    try:
+        result = subprocess.run(
+            ["security", "find-generic-password", "-s", "Claude Code-credentials", "-w"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        data = anyenv.load_json(result.stdout.strip(), return_type=dict)
+        val = data.get("claudeAiOauth", {}).get("accessToken")
+    except (subprocess.CalledProcessError, anyenv.JsonLoadError, FileNotFoundError):
+        return None
+    else:
+        assert isinstance(val, str)
+        return val
+
+
+def get_access_token() -> str | None:
+    """Get Claude Code OAuth access token.
+
+    Checks both file-based storage (Linux/newer macOS) and Keychain (macOS).
+
+    Returns:
+        Access token string, or None if not found.
+    """
+    # Try file-based first (works on Linux and newer macOS)
+    creds_path = _get_credentials_path()
+    if creds_path:
+        token = _get_access_token_from_file(creds_path)
+        if token:
+            return token
+    # Fall back to macOS Keychain
+    return _get_access_token_from_keychain()
+
+
+async def get_usage_async(token: str | None = None) -> ClaudeCodeUsage:
+    """Fetch Claude Code usage limits asynchronously.
+
+    Args:
+        token: OAuth access token. If not provided, will attempt to read from
+               stored credentials.
+
+    Returns:
+        ClaudeCodeUsage with current limits.
+
+    Raises:
+        ValueError: If no token provided and credentials not found.
+        httpx.HTTPStatusError: If API request fails.
+    """
+    if token is None:
+        token = get_access_token()
+        if token is None:
+            msg = "No Claude Code credentials found. Please authenticate with Claude Code first."
+            raise ValueError(msg)
+
+    async with httpx.AsyncClient() as client:
+        response = await client.get(
+            "https://api.anthropic.com/api/oauth/usage",
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+                "User-Agent": "claude-code/2.0.32",
+                "Authorization": f"Bearer {token}",
+                "anthropic-beta": "oauth-2025-04-20",
+            },
+        )
+        response.raise_for_status()
+        return ClaudeCodeUsage.model_validate(response.json())
+
+
+def get_usage(token: str | None = None) -> ClaudeCodeUsage:
+    """Fetch Claude Code usage limits synchronously.
+
+    Args:
+        token: OAuth access token. If not provided, will attempt to read from
+               stored credentials.
+
+    Returns:
+        ClaudeCodeUsage with current limits.
+
+    Raises:
+        ValueError: If no token provided and credentials not found.
+        httpx.HTTPStatusError: If API request fails.
+    """
+    if token is None:
+        token = get_access_token()
+        if token is None:
+            msg = "No Claude Code credentials found. Please authenticate with Claude Code first."
+            raise ValueError(msg)
+
+    with httpx.Client() as client:
+        response = client.get(
+            "https://api.anthropic.com/api/oauth/usage",
+            headers={
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+                "User-Agent": "claude-code/2.0.32",
+                "Authorization": f"Bearer {token}",
+                "anthropic-beta": "oauth-2025-04-20",
+            },
+        )
+        response.raise_for_status()
+        return ClaudeCodeUsage.model_validate(response.json())
+
+
+if __name__ == "__main__":
+    # Quick test
+    try:
+        usage = get_usage()
+        print(usage.format_table())
+    except ValueError as e:
+        print(f"Error: {e}")
+    except httpx.HTTPStatusError as e:
+        print(f"API Error: {e.response.status_code} - {e.response.text}")

agentpool/agents/context.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from contextvars import ContextVar
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Literal
 
@@ -9,6 +10,45 @@ from agentpool.log import get_logger
 from agentpool.messaging.context import NodeContext
 
 
+if TYPE_CHECKING:
+    from contextvars import Token
+
+
+# ContextVar for passing deps through async call boundaries (e.g., MCP tool bridge)
+# This allows run_stream() to set deps that are accessible in tool invocations
+_current_deps: ContextVar[Any] = ContextVar("current_deps", default=None)
+
+
+def set_current_deps(deps: Any) -> Token[Any]:
+    """Set the current deps for the running context.
+
+    Args:
+        deps: Dependencies to set
+
+    Returns:
+        Token to reset the deps when done
+    """
+    return _current_deps.set(deps)
+
+
+def get_current_deps() -> Any:
+    """Get the current deps from the running context.
+
+    Returns:
+        Current deps or None if not set
+    """
+    return _current_deps.get()
+
+
+def reset_current_deps(token: Token[Any]) -> None:
+    """Reset deps to previous value.
+
+    Args:
+        token: Token from set_current_deps
+    """
+    _current_deps.reset(token)
+
+
 if TYPE_CHECKING:
     from mcp import types
 

agentpool/agents/events/__init__.py

@@ -3,22 +3,27 @@
 from .events import (
     CommandCompleteEvent,
     CommandOutputEvent,
+    CompactionEvent,
     CustomEvent,
     DiffContentItem,
     FileContentItem,
     LocationContentItem,
     PlanUpdateEvent,
+    PartStartEvent,
+    PartDeltaEvent,
     RichAgentStreamEvent,
     RunErrorEvent,
     RunStartedEvent,
     SlashedAgentStreamEvent,
     StreamCompleteEvent,
+    SubAgentEvent,
     TerminalContentItem,
     TextContentItem,
     ToolCallCompleteEvent,
     ToolCallContentItem,
     ToolCallProgressEvent,
     ToolCallStartEvent,
+    ToolResultMetadataEvent,
 )
 from .event_emitter import StreamEventEmitter
 from .builtin_handlers import (
@@ -31,17 +36,30 @@ from .tts_handlers import (
     EdgeTTSEventHandler,
     OpenAITTSEventHandler,
 )
+from .processors import (
+    FileTracker,
+    FileTrackingProcessor,
+    StreamPipeline,
+    StreamProcessor,
+    event_handler_processor,
+    extract_file_path_from_tool_call,
+)
 
 __all__ = [
     "BaseTTSEventHandler",
     "CommandCompleteEvent",
     "CommandOutputEvent",
+    "CompactionEvent",
     "CustomEvent",
     "DiffContentItem",
     "EdgeTTSEventHandler",
     "FileContentItem",
+    "FileTracker",
+    "FileTrackingProcessor",
     "LocationContentItem",
     "OpenAITTSEventHandler",
+    "PartDeltaEvent",
+    "PartStartEvent",
     "PlanUpdateEvent",
     "RichAgentStreamEvent",
     "RunErrorEvent",
@@ -49,13 +67,19 @@ __all__ = [
     "SlashedAgentStreamEvent",
     "StreamCompleteEvent",
     "StreamEventEmitter",
+    "StreamPipeline",
+    "StreamProcessor",
+    "SubAgentEvent",
     "TerminalContentItem",
     "TextContentItem",
     "ToolCallCompleteEvent",
     "ToolCallContentItem",
     "ToolCallProgressEvent",
     "ToolCallStartEvent",
+    "ToolResultMetadataEvent",
     "detailed_print_handler",
+    "event_handler_processor",
+    "extract_file_path_from_tool_call",
     "resolve_event_handlers",
     "simple_print_handler",
 ]

agentpool/agents/events/builtin_handlers.py

@@ -24,6 +24,7 @@ from agentpool.agents.events import (
     ToolCallProgressEvent,
     ToolCallStartEvent,
 )
+from agentpool.utils.pydantic_ai_helpers import safe_args_as_dict
 
 
 if TYPE_CHECKING:
@@ -52,7 +53,7 @@ async def simple_print_handler(ctx: RunContext, event: RichAgentStreamEvent[Any]
             print(delta, end="", flush=True, file=sys.stderr)
 
         case FunctionToolCallEvent(part=ToolCallPart() as part):
-            kwargs_str = ", ".join(f"{k}={v!r}" for k, v in part.args_as_dict().items())
+            kwargs_str = ", ".join(f"{k}={v!r}" for k, v in safe_args_as_dict(part).items())
             print(f"\n🔧 {part.tool_name}({kwargs_str})", flush=True, file=sys.stderr)
 
         case FunctionToolResultEvent(result=ToolReturnPart() as return_part):
@@ -119,6 +120,71 @@ async def detailed_print_handler(ctx: RunContext, event: RichAgentStreamEvent[An
             print(file=sys.stderr)  # Final newline
 
 
+def create_file_stream_handler(
+    path: str,
+    mode: str = "a",
+    include_tools: bool = False,
+    include_thinking: bool = False,
+) -> IndividualEventHandler:
+    """Create an event handler that streams text output to a file.
+
+    Args:
+        path: Path to the output file
+        mode: File open mode ('w' for overwrite, 'a' for append)
+        include_tools: Whether to include tool call/result information
+        include_thinking: Whether to include thinking content
+
+    Returns:
+        Event handler function that writes to the specified file
+    """
+    from pathlib import Path
+
+    file_path = Path(path).expanduser()
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Open file handle that persists across calls
+    file_handle = file_path.open(mode, encoding="utf-8")
+
+    async def file_stream_handler(ctx: RunContext, event: RichAgentStreamEvent[Any]) -> None:
+        """Stream agent output to file."""
+        match event:
+            case (
+                PartStartEvent(part=TextPart(content=delta))
+                | PartDeltaEvent(delta=TextPartDelta(content_delta=delta))
+            ):
+                file_handle.write(delta)
+                file_handle.flush()
+
+            case (
+                PartStartEvent(part=ThinkingPart(content=delta))
+                | PartDeltaEvent(delta=ThinkingPartDelta(content_delta=delta))
+            ):
+                if include_thinking and delta:
+                    file_handle.write(f"\n[thinking] {delta}")
+                    file_handle.flush()
+
+            case FunctionToolCallEvent(part=ToolCallPart() as part):
+                if include_tools:
+                    kwargs_str = ", ".join(f"{k}={v!r}" for k, v in safe_args_as_dict(part).items())
+                    file_handle.write(f"\n[tool] {part.tool_name}({kwargs_str})\n")
+                    file_handle.flush()
+
+            case FunctionToolResultEvent(result=ToolReturnPart() as return_part):
+                if include_tools:
+                    file_handle.write(f"[result] {return_part.content}\n")
+                    file_handle.flush()
+
+            case RunErrorEvent(message=message):
+                file_handle.write(f"\n[error] {message}\n")
+                file_handle.flush()
+
+            case StreamCompleteEvent():
+                file_handle.write("\n")
+                file_handle.flush()
+
+    return file_stream_handler
+
+
 def resolve_event_handlers(
     event_handlers: Sequence[IndividualEventHandler | BuiltinEventHandlerType] | None,
 ) -> list[IndividualEventHandler]:

agentpool/agents/events/event_emitter.py

@@ -8,12 +8,15 @@ from agentpool.agents.events import (
     CustomEvent,
     LocationContentItem,
     PlanUpdateEvent,
+    TextContentItem,
     ToolCallProgressEvent,
     ToolCallStartEvent,
 )
 
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from agentpool.agents.context import AgentContext
     from agentpool.agents.events import RichAgentStreamEvent, ToolCallContentItem
     from agentpool.resource_providers.plan_provider import PlanEntry
@@ -78,7 +81,7 @@ class StreamEventEmitter:
         title: str,
         *,
         status: Literal["pending", "in_progress", "completed", "failed"] = "in_progress",
-        items: list[ToolCallContentItem] | None = None,
+        items: Sequence[ToolCallContentItem | str] | None = None,
         replace_content: bool = False,
     ) -> None:
         """Emit a progress event.
@@ -89,12 +92,26 @@ class StreamEventEmitter:
             items: Rich content items (terminals, diffs, locations, text)
             replace_content: If True, items replace existing content instead of appending
         """
+        # Record file changes from DiffContentItem for diff/revert support
+        if self._context.pool and self._context.pool.file_ops and items:
+            from agentpool.agents.events import DiffContentItem
+
+            for item in items:
+                if isinstance(item, DiffContentItem):
+                    self._context.pool.file_ops.record_change(
+                        path=item.path,
+                        old_content=item.old_text,
+                        new_content=item.new_text,
+                        operation="create" if item.old_text is None else "write",
+                        agent_name=self._context.node_name,
+                    )
+
         event = ToolCallProgressEvent(
             tool_call_id=self._context.tool_call_id or "",
             tool_name=self._context.tool_name,
             status=status,
             title=title,
-            items=items or [],
+            items=[TextContentItem(text=i) if isinstance(i, str) else i for i in items or []],
             replace_content=replace_content,
         )
         await self._emit(event)
@@ -129,6 +146,7 @@ class StreamEventEmitter:
         path: str,
         success: bool,
         error: str | None = None,
+        line: int = 0,
     ) -> None:
         """Emit file operation event.
 
@@ -137,6 +155,7 @@ class StreamEventEmitter:
             path: The file/directory path that was operated on
             success: Whether the operation completed successfully
             error: Error message if operation failed
+            line: Line number for navigation (0 = beginning)
         """
         event = ToolCallProgressEvent.file_operation(
             tool_call_id=self._context.tool_call_id or "",
@@ -145,6 +164,7 @@ class StreamEventEmitter:
             path=path,
             success=success,
             error=error,
+            line=line,
         )
         await self._emit(event)
 
@@ -163,6 +183,16 @@ class StreamEventEmitter:
             new_text: New file content
             status: Current status of the edit operation
         """
+        # Record file change for diff/revert support
+        if self._context.pool and self._context.pool.file_ops:
+            self._context.pool.file_ops.record_change(
+                path=path,
+                old_content=old_text,
+                new_content=new_text,
+                operation="edit",
+                agent_name=self._context.node_name,
+            )
+
         event = ToolCallProgressEvent.file_edit(
             tool_call_id=self._context.tool_call_id or "",
             tool_name=self._context.tool_name,