spatialagent/hooks_deep.py

"""
Hooks Middleware for deepagents integration.

Refactored hooks system that integrates with deepagents middleware architecture.
Supports both bash command hooks and prompt-based (LLM) hooks.
"""

import os
import json
import subprocess
import logging
import re
from typing import Dict, Any, List, Optional, Literal
from dataclasses import dataclass, field
from enum import Enum

from langchain.agents.middleware.types import AgentMiddleware
from langchain_core.messages import ToolMessage

logger = logging.getLogger(__name__)


class HookEvent(str, Enum):
    START = "Start"
    STOP = "Stop"
    PRE_PLAN = "PrePlan"
    POST_PLAN = "PostPlan"
    PRE_ACT = "PreAct"
    POST_ACT = "PostAct"
    PRE_TOOL_USE = "PreToolUse"
    POST_TOOL_USE = "PostToolUse"
    PRE_ROUTE = "PreRoute"
    PRE_CONCLUSION = "PreConclusion"


@dataclass
class HookResult:
    decision: Literal["approve", "block"] = "approve"
    reason: str = ""
    output: str = ""
    modified_args: Dict[str, Any] = field(default_factory=dict)

    @property
    def approved(self) -> bool:
        return self.decision == "approve"

    @property
    def blocked(self) -> bool:
        return self.decision == "block"


@dataclass
class HookDefinition:
    type: Literal["bash", "prompt"]
    command: Optional[str] = None
    prompt: Optional[str] = None
    timeout: int = 30
    matcher: Optional[Dict[str, Any]] = None


class HooksManager:
    def __init__(self, llm=None):
        self.llm = llm
        self.hooks: Dict[str, List[HookDefinition]] = {}
        self.enabled = True
        self._load_config()

    def _load_config(self):
        package_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        config_file = os.path.join(package_dir, ".spatialagent", "settings.json")

        if not os.path.exists(config_file):
            logger.debug("No .spatialagent/settings.json found, hooks disabled")
            self.enabled = False
            return

        try:
            with open(config_file, 'r') as f:
                config = json.load(f)

            hooks_config = config.get("hooks", {})

            for event_name, hook_list in hooks_config.items():
                try:
                    event = HookEvent(event_name)
                except ValueError:
                    logger.warning(f"Unknown hook event: {event_name}")
                    continue

                self.hooks[event_name] = []

                for hook_def in hook_list:
                    if "hooks" in hook_def:
                        for nested in hook_def["hooks"]:
                            self.hooks[event_name].append(self._parse_hook(nested, hook_def.get("matcher")))
                    else:
                        self.hooks[event_name].append(self._parse_hook(hook_def))

            logger.info(f"Loaded hooks from {config_file}: {list(self.hooks.keys())}")

        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON in {config_file}: {e}")
            self.enabled = False
        except Exception as e:
            logger.error(f"Error loading hooks config: {e}")
            self.enabled = False

    def _parse_hook(self, hook_def: Dict, parent_matcher: Optional[Dict] = None) -> HookDefinition:
        matcher = hook_def.get("matcher", parent_matcher)
        return HookDefinition(
            type=hook_def.get("type", "bash"),
            command=hook_def.get("command"),
            prompt=hook_def.get("prompt"),
            timeout=hook_def.get("timeout", 30),
            matcher=matcher
        )

    def set_llm(self, llm):
        self.llm = llm

    def _matches(self, hook: HookDefinition, context: Dict[str, Any]) -> bool:
        if not hook.matcher:
            return True

        for key, expected in hook.matcher.items():
            actual = context.get(key)
            if actual is None:
                return False

            if isinstance(expected, str) and isinstance(actual, str):
                if not re.match(expected, actual):
                    return False
            elif actual != expected:
                return False

        return True

    def _substitute_variables(self, template: str, context: Dict[str, Any]) -> str:
        result = template

        for key, value in context.items():
            placeholder = f"${key.upper()}"
            if placeholder in result:
                result = result.replace(placeholder, str(value))

        if "$ARGUMENTS" in result:
            result = result.replace("$ARGUMENTS", json.dumps(context, default=str))

        return result

    def _execute_bash_hook(self, hook: HookDefinition, context: Dict[str, Any]) -> HookResult:
        if not hook.command:
            return HookResult(decision="approve", reason="No command specified")

        command = self._substitute_variables(hook.command, context)

        try:
            env = os.environ.copy()
            for key, value in context.items():
                env[key.upper()] = str(value)

            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=hook.timeout,
                env=env
            )

            output = result.stdout + result.stderr

            try:
                decision_match = re.search(r'\{[^{}]*"decision"[^{}]*\}', output)
                if decision_match:
                    decision_json = json.loads(decision_match.group())
                    return HookResult(
                        decision=decision_json.get("decision", "approve"),
                        reason=decision_json.get("reason", ""),
                        output=output
                    )
            except json.JSONDecodeError:
                pass

            if result.returncode == 0:
                return HookResult(decision="approve", output=output)
            else:
                return HookResult(
                    decision="block",
                    reason=f"Command exited with code {result.returncode}",
                    output=output
                )

        except subprocess.TimeoutExpired:
            return HookResult(
                decision="block",
                reason=f"Hook timed out after {hook.timeout}s"
            )
        except Exception as e:
            logger.error(f"Bash hook error: {e}")
            return HookResult(decision="approve", reason=f"Hook error: {e}")

    def _execute_prompt_hook(self, hook: HookDefinition, context: Dict[str, Any]) -> HookResult:
        if not self.llm:
            logger.warning("Prompt hook requested but no LLM configured")
            return HookResult(decision="approve", reason="No LLM for prompt hook")

        if not hook.prompt:
            return HookResult(decision="approve", reason="No prompt specified")

        prompt = self._substitute_variables(hook.prompt, context)

        try:
            from langchain_core.messages import HumanMessage

            response = self.llm.invoke([HumanMessage(content=prompt)])
            output = str(response.content)

            try:
                json_match = re.search(r'\{[^{}]*"decision"[^{}]*\}', output, re.DOTALL)
                if json_match:
                    decision_json = json.loads(json_match.group())
                    return HookResult(
                        decision=decision_json.get("decision", "approve"),
                        reason=decision_json.get("reason", ""),
                        output=output,
                        modified_args=decision_json.get("modified_args", {})
                    )
            except json.JSONDecodeError:
                pass

            output_lower = output.lower()
            if "block" in output_lower or "deny" in output_lower or "reject" in output_lower:
                return HookResult(decision="block", reason=output, output=output)

            return HookResult(decision="approve", output=output)

        except Exception as e:
            logger.error(f"Prompt hook error: {e}")
            return HookResult(decision="approve", reason=f"Hook error: {e}")

    def execute(self, event: HookEvent, context: Dict[str, Any]) -> HookResult:
        if not self.enabled:
            return HookResult(decision="approve")

        event_name = event.value if isinstance(event, HookEvent) else event
        hooks = self.hooks.get(event_name, [])

        if not hooks:
            return HookResult(decision="approve")

        logger.debug(f"Executing {len(hooks)} hooks for {event_name}")

        all_outputs = []
        modified_args = {}

        for hook in hooks:
            if not self._matches(hook, context):
                continue

            if hook.type == "bash":
                result = self._execute_bash_hook(hook, context)
            elif hook.type == "prompt":
                result = self._execute_prompt_hook(hook, context)
            else:
                logger.warning(f"Unknown hook type: {hook.type}")
                continue

            all_outputs.append(result.output)
            modified_args.update(result.modified_args)

            if result.blocked:
                logger.info(f"Hook blocked {event_name}: {result.reason}")
                return HookResult(
                    decision="block",
                    reason=result.reason,
                    output="\n".join(all_outputs),
                    modified_args=modified_args
                )

        return HookResult(
            decision="approve",
            output="\n".join(all_outputs),
            modified_args=modified_args
        )

    def has_hooks(self, event: HookEvent) -> bool:
        event_name = event.value if isinstance(event, HookEvent) else event
        return bool(self.hooks.get(event_name))


class HooksMiddleware(AgentMiddleware):
    """
    Deepagents middleware that integrates the hooks system.

    Maps deepagents lifecycle events to hook events:
    - before_agent -> HookEvent.START
    - after_agent -> HookEvent.STOP
    - before_model -> HookEvent.PRE_PLAN
    - after_model -> HookEvent.POST_PLAN
    - wrap_tool_call -> HookEvent.PRE_TOOL_USE / HookEvent.POST_TOOL_USE
    """

    def __init__(self, llm=None):
        self.hooks_manager = HooksManager(llm=llm)

    def before_agent(self, state, runtime):
        """Run before agent execution starts."""
        context = {
            "state": str(state),
            "runtime": str(runtime),
        }
        result = self.hooks_manager.execute(HookEvent.START, context)
        if result.blocked:
            logger.warning(f"Hook blocked START event: {result.reason}")
        return None

    def after_agent(self, state, runtime):
        """Run after agent execution completes."""
        context = {
            "state": str(state),
            "runtime": str(runtime),
        }
        result = self.hooks_manager.execute(HookEvent.STOP, context)
        if result.blocked:
            logger.warning(f"Hook blocked STOP event: {result.reason}")
        return None

    def before_model(self, state, runtime):
        """Run before model is called (planning phase)."""
        context = {
            "state": str(state),
            "runtime": str(runtime),
        }
        result = self.hooks_manager.execute(HookEvent.PRE_PLAN, context)
        if result.blocked:
            logger.warning(f"Hook blocked PRE_PLAN event: {result.reason}")
        return None

    def after_model(self, state, runtime):
        """Run after model is called (planning phase complete)."""
        context = {
            "state": str(state),
            "runtime": str(runtime),
        }
        result = self.hooks_manager.execute(HookEvent.POST_PLAN, context)
        if result.blocked:
            logger.warning(f"Hook blocked POST_PLAN event: {result.reason}")
        return None

    def wrap_tool_call(self, request, handler):
        """Intercept tool execution for hooks."""
        tool_name = request.tool_call.get("name", "")
        tool_args = request.tool_call.get("args", {})

        pre_context = {
            "tool_name": tool_name,
            "tool_args": json.dumps(tool_args),
            "state": str(request.state),
            "runtime": str(request.runtime),
        }

        pre_result = self.hooks_manager.execute(HookEvent.PRE_TOOL_USE, pre_context)
        if pre_result.blocked:
            logger.warning(f"Hook blocked PRE_TOOL_USE for {tool_name}: {pre_result.reason}")
            return ToolMessage(
                content=f"Tool call blocked by hook: {pre_result.reason}",
                tool_call_id=request.tool_call.get("id"),
                status="error"
            )

        if pre_result.modified_args:
            modified_call = {
                **request.tool_call,
                "args": {**tool_args, **pre_result.modified_args}
            }
            request = request.override(tool_call=modified_call)
            tool_args = modified_call["args"]

        try:
            result = handler(request)

            post_context = {
                "tool_name": tool_name,
                "tool_args": json.dumps(tool_args),
                "tool_result": str(result),
                "state": str(request.state),
                "runtime": str(request.runtime),
            }

            post_result = self.hooks_manager.execute(HookEvent.POST_TOOL_USE, post_context)
            if post_result.blocked:
                logger.warning(f"Hook blocked POST_TOOL_USE for {tool_name}: {post_result.reason}")

            return result
        except Exception as e:
            error_context = {
                "tool_name": tool_name,
                "tool_args": json.dumps(tool_args),
                "error": str(e),
                "state": str(request.state),
                "runtime": str(request.runtime),
            }
            self.hooks_manager.execute(HookEvent.POST_TOOL_USE, error_context)
            raise


# Global hooks manager instance
_hooks_manager: Optional[HooksManager] = None


def get_hooks_manager() -> HooksManager:
    global _hooks_manager
    if _hooks_manager is None:
        _hooks_manager = HooksManager()
    return _hooks_manager


def set_hooks_manager(manager: HooksManager):
    global _hooks_manager
    _hooks_manager = manager


def init_hooks(llm=None) -> HooksManager:
    manager = HooksManager(llm=llm)
    set_hooks_manager(manager)
    return manager