agentpool 2.1.9__py3-none-any.whl

agentpool/hooks/agent_hooks.py

@@ -0,0 +1,247 @@
+"""AgentHooks - Runtime hook container for agent lifecycle events."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from agentpool.hooks.base import HookResult
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from agentpool.hooks.base import Hook, HookInput
+
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class AgentHooks:
+    """Runtime container for agent lifecycle hooks.
+
+    Holds instantiated hooks organized by event type and provides
+    methods to execute them with proper input/output handling.
+
+    Attributes:
+        pre_run: Hooks executed before agent.run() processes a prompt.
+        post_run: Hooks executed after agent.run() completes.
+        pre_tool_use: Hooks executed before a tool is called.
+        post_tool_use: Hooks executed after a tool completes.
+    """
+
+    pre_run: list[Hook] = field(default_factory=list)
+    post_run: list[Hook] = field(default_factory=list)
+    pre_tool_use: list[Hook] = field(default_factory=list)
+    post_tool_use: list[Hook] = field(default_factory=list)
+
+    def has_hooks(self) -> bool:
+        """Check if any hooks are configured."""
+        return bool(self.pre_run or self.post_run or self.pre_tool_use or self.post_tool_use)
+
+    async def run_pre_run_hooks(
+        self,
+        *,
+        agent_name: str,
+        prompt: str,
+        conversation_id: str | None = None,
+    ) -> HookResult:
+        """Execute pre-run hooks.
+
+        Args:
+            agent_name: Name of the agent.
+            prompt: The prompt being processed.
+            conversation_id: Optional conversation identifier.
+
+        Returns:
+            Combined hook result. If any hook denies, the run should be blocked.
+        """
+        input_data: HookInput = {
+            "event": "pre_run",
+            "agent_name": agent_name,
+            "prompt": prompt,
+            "conversation_id": conversation_id,
+        }
+        return await self._run_hooks(self.pre_run, input_data)
+
+    async def run_post_run_hooks(
+        self,
+        *,
+        agent_name: str,
+        prompt: str,
+        result: Any,
+        conversation_id: str | None = None,
+    ) -> HookResult:
+        """Execute post-run hooks.
+
+        Args:
+            agent_name: Name of the agent.
+            prompt: The prompt that was processed.
+            result: The result from the run.
+            conversation_id: Optional conversation identifier.
+
+        Returns:
+            Combined hook result.
+        """
+        input_data: HookInput = {
+            "event": "post_run",
+            "agent_name": agent_name,
+            "prompt": prompt,
+            "result": result,
+            "conversation_id": conversation_id,
+        }
+        return await self._run_hooks(self.post_run, input_data)
+
+    async def run_pre_tool_hooks(
+        self,
+        *,
+        agent_name: str,
+        tool_name: str,
+        tool_input: dict[str, Any],
+        conversation_id: str | None = None,
+    ) -> HookResult:
+        """Execute pre-tool-use hooks.
+
+        Args:
+            agent_name: Name of the agent.
+            tool_name: Name of the tool being called.
+            tool_input: Input arguments for the tool.
+            conversation_id: Optional conversation identifier.
+
+        Returns:
+            Combined hook result. If any hook denies, the tool call should be blocked.
+            May include modified_input to change tool arguments.
+        """
+        input_data: HookInput = {
+            "event": "pre_tool_use",
+            "agent_name": agent_name,
+            "tool_name": tool_name,
+            "tool_input": tool_input,
+            "conversation_id": conversation_id,
+        }
+        return await self._run_hooks(self.pre_tool_use, input_data)
+
+    async def run_post_tool_hooks(
+        self,
+        *,
+        agent_name: str,
+        tool_name: str,
+        tool_input: dict[str, Any],
+        tool_output: Any,
+        duration_ms: float,
+        conversation_id: str | None = None,
+    ) -> HookResult:
+        """Execute post-tool-use hooks.
+
+        Args:
+            agent_name: Name of the agent.
+            tool_name: Name of the tool that was called.
+            tool_input: Input arguments that were passed to the tool.
+            tool_output: Output from the tool.
+            duration_ms: How long the tool took to execute.
+            conversation_id: Optional conversation identifier.
+
+        Returns:
+            Combined hook result. May include additional_context to inject.
+        """
+        input_data: HookInput = {
+            "event": "post_tool_use",
+            "agent_name": agent_name,
+            "tool_name": tool_name,
+            "tool_input": tool_input,
+            "tool_output": tool_output,
+            "duration_ms": duration_ms,
+            "conversation_id": conversation_id,
+        }
+        return await self._run_hooks(self.post_tool_use, input_data)
+
+    @staticmethod
+    async def _run_hooks(hooks: list[Hook], input_data: HookInput) -> HookResult:
+        """Run a list of hooks and combine their results.
+
+        Hooks are run in parallel. Results are combined:
+        - If any hook returns "deny", the combined result is "deny"
+        - If any hook returns "ask", the combined result is "ask" (unless denied)
+        - Reasons are concatenated
+        - modified_input values are merged (later hooks override earlier)
+        - additional_context values are concatenated
+        - continue_ is False if any hook sets it False
+
+        Args:
+            hooks: List of hooks to execute.
+            input_data: Input data for the hooks.
+
+        Returns:
+            Combined hook result.
+        """
+        if not hooks:
+            return HookResult(decision="allow")
+
+        # Filter to matching hooks
+        matching = [h for h in hooks if h.matches(input_data)]
+        if not matching:
+            return HookResult(decision="allow")
+
+        # Run all matching hooks in parallel
+        raw_results = await asyncio.gather(
+            *(hook.execute(input_data) for hook in matching),
+            return_exceptions=True,
+        )
+
+        # Combine results
+        combined: HookResult = {"decision": "allow"}
+        reasons: list[str] = []
+        contexts: list[str] = []
+
+        for raw_result in raw_results:
+            if isinstance(raw_result, BaseException):
+                logger.warning("Hook execution failed", error=str(raw_result))
+                continue
+
+            result: HookResult = raw_result
+
+            # Decision priority: deny > ask > allow
+            if result.get("decision") == "deny":
+                combined["decision"] = "deny"
+            elif result.get("decision") == "ask" and combined.get("decision") != "deny":
+                combined["decision"] = "ask"
+
+            # Collect reasons
+            if reason := result.get("reason"):
+                reasons.append(reason)
+
+            # Merge modified_input (later overrides earlier)
+            if modified := result.get("modified_input"):
+                if "modified_input" not in combined:
+                    combined["modified_input"] = {}
+                combined["modified_input"].update(modified)
+
+            # Collect additional context
+            if ctx := result.get("additional_context"):
+                contexts.append(ctx)
+
+            # continue_ is False if any hook sets it False
+            if result.get("continue_") is False:
+                combined["continue_"] = False
+
+        # Combine collected values
+        if reasons:
+            combined["reason"] = "; ".join(reasons)
+        if contexts:
+            combined["additional_context"] = "\n".join(contexts)
+
+        return combined
+
+    def __repr__(self) -> str:
+        counts = {
+            "pre_run": len(self.pre_run),
+            "post_run": len(self.post_run),
+            "pre_tool_use": len(self.pre_tool_use),
+            "post_tool_use": len(self.post_tool_use),
+        }
+        non_empty = {k: v for k, v in counts.items() if v > 0}
+        if not non_empty:
+            return "AgentHooks(empty)"
+        parts = ", ".join(f"{k}={v}" for k, v in non_empty.items())
+        return f"AgentHooks({parts})"

agentpool/hooks/base.py

@@ -0,0 +1,119 @@
+"""Base hook classes and types."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+import re
+from typing import Any, Literal, TypedDict
+
+from agentpool.log import get_logger
+
+
+logger = get_logger(__name__)
+
+HookEvent = Literal["pre_run", "post_run", "pre_tool_use", "post_tool_use"]
+
+
+class HookInput(TypedDict, total=False):
+    """Input data passed to hooks."""
+
+    # Common fields
+    event: HookEvent
+    agent_name: str
+    conversation_id: str | None
+
+    # Tool-related fields (pre_tool_use, post_tool_use)
+    tool_name: str
+    tool_input: dict[str, Any]
+    tool_output: Any
+    duration_ms: float
+
+    # Run-related fields (pre_run, post_run)
+    prompt: str
+    result: Any
+
+
+class HookResult(TypedDict, total=False):
+    """Result returned from hook execution."""
+
+    decision: Literal["allow", "deny", "ask"]
+    """Decision for pre_* hooks: allow, deny, or ask user."""
+
+    reason: str
+    """Explanation for the decision."""
+
+    modified_input: dict[str, Any]
+    """Modified input for pre_* hooks (e.g., modified tool_input)."""
+
+    additional_context: str
+    """Context to inject into conversation."""
+
+    continue_: bool
+    """Whether to continue execution. False = stop."""
+
+
+class Hook(ABC):
+    """Base class for runtime hooks."""
+
+    def __init__(
+        self,
+        event: HookEvent,
+        matcher: str | None = None,
+        timeout: float = 60.0,
+        enabled: bool = True,
+    ) -> None:
+        """Initialize hook.
+
+        Args:
+            event: The lifecycle event this hook handles.
+            matcher: Regex pattern for matching (e.g., tool names). None matches all.
+            timeout: Maximum execution time in seconds.
+            enabled: Whether this hook is active.
+        """
+        self.event = event
+        self.matcher = matcher
+        self.timeout = timeout
+        self.enabled = enabled
+        self._pattern = re.compile(matcher) if matcher and matcher != "*" else None
+
+    def matches(self, input_data: HookInput) -> bool:
+        """Check if this hook should run for the given input.
+
+        Args:
+            input_data: The hook input data.
+
+        Returns:
+            True if the hook should execute.
+        """
+        if not self.enabled:
+            return False
+
+        # No pattern means match all
+        if self._pattern is None:
+            return True
+
+        # For tool events, match against tool_name
+        if self.event in ("pre_tool_use", "post_tool_use"):
+            tool_name = input_data.get("tool_name", "")
+            return bool(self._pattern.search(tool_name))
+
+        # For run events, pattern matching doesn't apply (no tool name to match)
+        return True
+
+    @abstractmethod
+    async def execute(self, input_data: HookInput) -> HookResult:
+        """Execute the hook.
+
+        Args:
+            input_data: The hook input data.
+
+        Returns:
+            Hook result with decision and optional modifications.
+        """
+        ...
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"event={self.event!r}, matcher={self.matcher!r}, enabled={self.enabled})"
+        )

agentpool/hooks/callable.py

@@ -0,0 +1,140 @@
+"""Callable hook implementation."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from agentpool.hooks.base import Hook, HookResult
+from agentpool.log import get_logger
+from agentpool.utils.importing import import_callable
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from agentpool.hooks.base import HookEvent, HookInput
+
+
+logger = get_logger(__name__)
+
+
+class CallableHook(Hook):
+    """Hook that executes a Python callable.
+
+    The callable receives hook input as a dictionary and should return
+    a HookResult dictionary or None.
+    """
+
+    def __init__(
+        self,
+        event: HookEvent,
+        fn: Callable[..., HookResult | None] | str,
+        matcher: str | None = None,
+        timeout: float = 60.0,
+        enabled: bool = True,
+        arguments: dict[str, Any] | None = None,
+    ):
+        """Initialize callable hook.
+
+        Args:
+            event: The lifecycle event this hook handles.
+            fn: The callable to execute, or import path string.
+            matcher: Regex pattern for matching.
+            timeout: Maximum execution time in seconds.
+            enabled: Whether this hook is active.
+            arguments: Additional keyword arguments for the callable.
+        """
+        super().__init__(event=event, matcher=matcher, timeout=timeout, enabled=enabled)
+        self._callable: Callable[..., HookResult | None] | None = None
+        self._import_path: str | None = None
+
+        if isinstance(fn, str):
+            self._import_path = fn
+        else:
+            self._callable = fn
+
+        self.arguments = arguments or {}
+
+    @property
+    def callable(self) -> Callable[..., HookResult | None]:
+        """Get the callable, importing lazily if needed."""
+        if self._callable is None:
+            if self._import_path is None:
+                msg = "No callable or import path provided"
+                raise ValueError(msg)
+            self._callable = import_callable(self._import_path)
+        return self._callable
+
+    async def execute(self, input_data: HookInput) -> HookResult:
+        """Execute the callable.
+
+        Args:
+            input_data: The hook input data.
+
+        Returns:
+            Hook result from callable.
+        """
+        try:
+            fn = self.callable
+            # Merge input data with additional arguments
+            kwargs = {**dict(input_data), **self.arguments}
+            # Execute with timeout
+            if asyncio.iscoroutinefunction(fn):
+                result = await asyncio.wait_for(fn(**kwargs), timeout=self.timeout)  # ty: ignore
+            else:
+                # Run sync function in executor
+                loop = asyncio.get_event_loop()
+                result = await asyncio.wait_for(
+                    loop.run_in_executor(None, lambda: fn(**kwargs)),
+                    timeout=self.timeout,
+                )
+
+            # Normalize result
+            if result is None:
+                return HookResult(decision="allow")
+
+            return _normalize_result(result)
+
+        except TimeoutError:
+            fn_path = self._import_path or str(self._callable)
+            logger.exception("Hook callable timed out", timeout=self.timeout, callable=fn_path)
+            return HookResult(decision="allow")
+        except Exception as e:
+            fn_path = self._import_path or str(self._callable)
+            logger.exception("Hook callable failed", callable=fn_path)
+            return HookResult(decision="allow", reason=str(e))
+
+
+def _normalize_result(result: Any) -> HookResult:
+    """Normalize callable result to HookResult.
+
+    Args:
+        result: Result from callable.
+
+    Returns:
+        Normalized hook result.
+    """
+    if isinstance(result, dict):
+        # Already a dict, ensure proper typing
+        normalized: HookResult = {}
+        if "decision" in result:
+            normalized["decision"] = result["decision"]
+        if "reason" in result:
+            normalized["reason"] = result["reason"]
+        if "modified_input" in result:
+            normalized["modified_input"] = result["modified_input"]
+        if "additional_context" in result:
+            normalized["additional_context"] = result["additional_context"]
+        if "continue_" in result:
+            normalized["continue_"] = result["continue_"]
+        return normalized
+
+    # String result treated as additional context
+    if isinstance(result, str):
+        return HookResult(decision="allow", additional_context=result)
+    # Bool result treated as allow/deny
+    if isinstance(result, bool):
+        return HookResult(decision="allow" if result else "deny")
+    # Unknown type, allow by default
+    return HookResult(decision="allow")

agentpool/hooks/command.py

@@ -0,0 +1,180 @@
+"""Command hook implementation."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from agentpool.hooks.base import Hook, HookResult
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from agentpool.hooks.base import HookEvent, HookInput
+
+
+logger = get_logger(__name__)
+
+
+class CommandHook(Hook):
+    """Hook that executes a shell command.
+
+    The command receives hook input as JSON via stdin and should return
+    JSON output via stdout.
+
+    Exit codes:
+    - 0: Success, stdout parsed as JSON for result
+    - 2: Block/deny, stderr used as reason
+    - Other: Non-blocking error, logged but execution continues
+    """
+
+    def __init__(
+        self,
+        event: HookEvent,
+        command: str,
+        matcher: str | None = None,
+        timeout: float = 60.0,
+        enabled: bool = True,
+        env: dict[str, str] | None = None,
+    ):
+        """Initialize command hook.
+
+        Args:
+            event: The lifecycle event this hook handles.
+            command: Shell command to execute.
+            matcher: Regex pattern for matching.
+            timeout: Maximum execution time in seconds.
+            enabled: Whether this hook is active.
+            env: Additional environment variables.
+        """
+        super().__init__(event=event, matcher=matcher, timeout=timeout, enabled=enabled)
+        self.command = command
+        self.env = env or {}
+
+    async def execute(self, input_data: HookInput) -> HookResult:
+        """Execute the shell command.
+
+        Args:
+            input_data: The hook input data, passed as JSON to stdin.
+
+        Returns:
+            Hook result parsed from command output.
+        """
+        # Prepare environment
+        env = os.environ.copy()
+        env.update(self.env)
+
+        # Expand $PROJECT_DIR if present
+        command = self.command
+        if "$PROJECT_DIR" in command:
+            project_dir = env.get("PROJECT_DIR", Path.cwd())
+            command = command.replace("$PROJECT_DIR", str(project_dir))
+
+        # Serialize input
+        input_json = json.dumps(dict(input_data))
+
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                command,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                env=env,
+            )
+
+            stdout, stderr = await asyncio.wait_for(
+                proc.communicate(input_json.encode()),
+                timeout=self.timeout,
+            )
+
+            stdout_str = stdout.decode().strip()
+            stderr_str = stderr.decode().strip()
+
+            # Handle exit codes
+            if proc.returncode == 0:
+                return _parse_success_output(stdout_str)
+            if proc.returncode == 2:  # noqa: PLR2004
+                # Blocking error
+                reason = stderr_str or "Hook denied the operation"
+                return HookResult(decision="deny", reason=reason)
+            # Non-blocking error
+            logger.warning("Hook command failed", returncode=proc.returncode, stderr=stderr_str)
+            return HookResult(decision="allow")
+
+        except TimeoutError:
+            logger.exception("Hook command timed out", timeout=self.timeout, command=command)
+            return HookResult(decision="allow")
+        except Exception as e:
+            logger.exception("Hook command failed", command=command)
+            return HookResult(decision="allow", reason=str(e))
+
+
+def _parse_success_output(stdout: str) -> HookResult:
+    """Parse successful command output.
+
+    Args:
+        stdout: Command stdout.
+
+    Returns:
+        Parsed hook result.
+    """
+    if not stdout:
+        return HookResult(decision="allow")
+
+    try:
+        data = json.loads(stdout)
+        return _normalize_result(data)
+    except json.JSONDecodeError:
+        # Plain text output treated as additional context
+        return HookResult(decision="allow", additional_context=stdout)
+
+
+def _normalize_result(data: dict[str, Any]) -> HookResult:
+    """Normalize command output to HookResult.
+
+    Args:
+        data: Parsed JSON data.
+
+    Returns:
+        Normalized hook result.
+    """
+    result: HookResult = {}
+
+    # Handle decision field (support various naming conventions)
+    decision = data.get("decision") or data.get("permissionDecision")
+    if decision:
+        # Normalize decision values
+        if decision in ("approve", "allow"):
+            result["decision"] = "allow"
+        elif decision in ("block", "deny"):
+            result["decision"] = "deny"
+        elif decision == "ask":
+            result["decision"] = "ask"
+
+    # Handle reason field
+    reason = data.get("reason") or data.get("permissionDecisionReason")
+    if reason:
+        result["reason"] = reason
+
+    # Handle modified input
+    if "modified_input" in data:
+        result["modified_input"] = data["modified_input"]
+    elif "updatedInput" in data:
+        result["modified_input"] = data["updatedInput"]
+
+    # Handle additional context
+    if "additional_context" in data:
+        result["additional_context"] = data["additional_context"]
+    elif "additionalContext" in data:
+        result["additional_context"] = data["additionalContext"]
+
+    # Handle continue flag
+    if "continue" in data:
+        result["continue_"] = data["continue"]
+    elif "continue_" in data:
+        result["continue_"] = data["continue_"]
+
+    return result