crewAI/lib/crewai/tests/hooks/test_interception_conformance.py

"""Conformance suite for the framework-native interception points.

For each wired point this suite asserts the shared contract: the probe hook
sees a well-shaped payload, an in-place/returned modification is honored, and a
:class:`HookAborted` interrupts the step.
"""

from __future__ import annotations

from crewai.flow.flow import Flow, listen, start
from crewai.hooks.dispatch import (
    HookAborted,
    InterceptionPoint,
    clear_all,
    on,
)
import pytest


@pytest.fixture(autouse=True)
def clear_dispatch_registry():
    clear_all()
    yield
    clear_all()


class _SimpleFlow(Flow):
    @start()
    def begin(self):
        return "begin"

    @listen(begin)
    def finish(self, _):
        return "flow-result"


class TestFlowExecutionBoundaries:
    """execution_start / input / output / execution_end on a flow."""

    def test_all_boundary_points_fire_once(self):
        fired: list[str] = []

        for point in (
            InterceptionPoint.EXECUTION_START,
            InterceptionPoint.INPUT,
            InterceptionPoint.OUTPUT,
            InterceptionPoint.EXECUTION_END,
        ):

            @on(point)
            def _probe(ctx, _point=point):
                fired.append(_point.value)

        _SimpleFlow().kickoff(inputs={"seed": 1})

        assert fired == [
            "execution_start",
            "input",
            "output",
            "execution_end",
        ]

    def test_output_modification_is_honored(self):
        @on(InterceptionPoint.OUTPUT)
        def rewrite(ctx):
            return "intercepted"

        result = _SimpleFlow().kickoff()
        assert result == "intercepted"

    def test_input_payload_carries_inputs(self):
        seen: dict = {}

        @on(InterceptionPoint.INPUT)
        def capture(ctx):
            seen.update(ctx.payload or {})

        _SimpleFlow().kickoff(inputs={"seed": 42})
        assert seen == {"seed": 42}

    def test_abort_at_execution_start_interrupts(self):
        @on(InterceptionPoint.EXECUTION_START)
        def block(ctx):
            raise HookAborted(reason="not allowed", source="policy")

        with pytest.raises(HookAborted) as exc:
            _SimpleFlow().kickoff()
        assert exc.value.reason == "not allowed"


class TestFlowStepPoints:
    """pre_step / post_step for flow methods (kind=flow_method)."""

    def test_pre_and_post_step_fire_per_method(self):
        kinds: list[tuple[str, str | None]] = []

        @on(InterceptionPoint.PRE_STEP)
        def pre(ctx):
            kinds.append(("pre", ctx.step_name))

        @on(InterceptionPoint.POST_STEP)
        def post(ctx):
            kinds.append(("post", ctx.step_name))

        _SimpleFlow().kickoff()

        assert ("pre", "begin") in kinds
        assert ("post", "begin") in kinds
        assert ("pre", "finish") in kinds
        assert ("post", "finish") in kinds

    def test_post_step_can_rewrite_method_output(self):
        @on(InterceptionPoint.POST_STEP)
        def rewrite(ctx):
            if ctx.step_name == "finish":
                return "rewritten"
            return None

        assert _SimpleFlow().kickoff() == "rewritten"


class TestToolSelection:
    """tool_selection must gate every path that hands tools to an executor."""

    @staticmethod
    def _make_agent():
        from crewai import Agent

        return Agent(
            role="Selector",
            goal="g",
            backstory="b",
            llm="gpt-4o-mini",
        )

    @staticmethod
    def _make_tools():
        from crewai.tools import tool

        @tool("allowed_tool")
        def allowed_tool() -> str:
            """A tool the policy permits."""
            return "ok"

        @tool("blocked_tool")
        def blocked_tool() -> str:
            """A tool the policy strips."""
            return "nope"

        return [allowed_tool, blocked_tool]

    def test_hook_filters_tools_on_executor_creation(self):
        @on(InterceptionPoint.TOOL_SELECTION)
        def drop_blocked(ctx):
            return [t for t in ctx.payload if t.name != "blocked_tool"]

        agent = self._make_agent()
        agent.create_agent_executor(tools=self._make_tools())

        names = [t.name for t in agent.agent_executor.original_tools]
        assert "allowed_tool" in names
        assert "blocked_tool" not in names

    def test_hook_filters_tools_on_standalone_kickoff_prepare(self):
        @on(InterceptionPoint.TOOL_SELECTION)
        def drop_blocked(ctx):
            return [t for t in ctx.payload if t.name != "blocked_tool"]

        agent = self._make_agent()
        agent.tools = self._make_tools()
        _executor, _inputs, agent_info, parsed_tools = agent._prepare_kickoff("hi")

        names = [t.name for t in agent_info["tools"]]
        assert "allowed_tool" in names
        assert "blocked_tool" not in names
        assert all(t.name != "blocked_tool" for t in parsed_tools)