vibe/tests/agent_loop/e2e/test_e2e_hooks.py

from __future__ import annotations

import json
import shlex
import sys
import textwrap
from typing import Any

import pytest

from tests.agent_loop.e2e.conftest import MistralAPI, build_e2e_agent_loop
from tests.backend.data.mistral import mistral_completion
from tests.conftest import build_test_vibe_config
from vibe.core.hooks.config import HookConfigResult
from vibe.core.hooks.models import HookConfig, HookType
from vibe.core.types import AssistantEvent, ToolResultEvent

TODO_TOOL_CALL: list[dict[str, Any]] = [
    {
        "id": "call_todo",
        "function": {"name": "todo", "arguments": '{"action": "read"}'},
        "index": 0,
    }
]


def _emit_cmd(payload: dict[str, Any]) -> str:
    # Payload is passed via stdin of the runner (stdout of parent process)
    body = f"import sys; sys.stdout.write({json.dumps(payload)!r})"
    return f"{sys.executable} -c {shlex.quote(body)}"


@pytest.mark.asyncio
async def test_before_tool_hook_denies_tool_and_skips_execution(
    mistral_api: MistralAPI,
) -> None:
    deny_hook = HookConfig(
        name="deny",
        type=HookType.BEFORE_TOOL,
        command=_emit_cmd({"decision": "deny", "reason": "blocked by policy"}),
        match="todo",
    )
    mistral_api.reply(
        mistral_completion("", tool_calls=TODO_TOOL_CALL), mistral_completion("done")
    )
    agent = build_e2e_agent_loop(
        config=build_test_vibe_config(enabled_tools=["todo"]),
        hook_config_result=HookConfigResult(hooks=[deny_hook], issues=[]),
    )

    events = [event async for event in agent.act("Show my todos")]

    tool_result = next(e for e in events if isinstance(e, ToolResultEvent))
    assert tool_result.skipped is True
    assert "blocked by policy" in (tool_result.skip_reason or "")


@pytest.mark.asyncio
async def test_before_tool_hook_rewrites_tool_input_seen_by_model(
    mistral_api: MistralAPI,
) -> None:
    rewritten = {
        "action": "write",
        "todos": [{"id": "1", "content": "rewritten by hook"}],
    }
    rewrite_hook = HookConfig(
        name="rewrite",
        type=HookType.BEFORE_TOOL,
        command=_emit_cmd({"hook_specific_output": {"tool_input": rewritten}}),
        match="todo",
    )
    mistral_api.reply(
        mistral_completion("", tool_calls=TODO_TOOL_CALL), mistral_completion("done")
    )
    agent = build_e2e_agent_loop(
        config=build_test_vibe_config(enabled_tools=["todo"]),
        hook_config_result=HookConfigResult(hooks=[rewrite_hook], issues=[]),
    )

    [_ async for _ in agent.act("Update my todos")]

    assert "rewritten by hook" in mistral_api.model_facing_text(1)


@pytest.mark.asyncio
async def test_before_tool_rewrite_failing_validation_is_denied(
    mistral_api: MistralAPI,
) -> None:
    bad_rewrite_hook = HookConfig(
        name="bad-rewrite",
        type=HookType.BEFORE_TOOL,
        command=_emit_cmd({"hook_specific_output": {"tool_input": {"todos": "x"}}}),
        match="todo",
    )
    mistral_api.reply(
        mistral_completion("", tool_calls=TODO_TOOL_CALL), mistral_completion("done")
    )
    agent = build_e2e_agent_loop(
        config=build_test_vibe_config(enabled_tools=["todo"]),
        hook_config_result=HookConfigResult(hooks=[bad_rewrite_hook], issues=[]),
    )

    events = [event async for event in agent.act("Update my todos")]

    tool_result = next(e for e in events if isinstance(e, ToolResultEvent))
    assert tool_result.skipped is True
    assert "failed validation" in (tool_result.skip_reason or "")


@pytest.mark.asyncio
async def test_after_tool_hook_replaces_tool_output_seen_by_model(
    mistral_api: MistralAPI,
) -> None:
    replace_output_hook = HookConfig(
        name="replace",
        type=HookType.AFTER_TOOL,
        command=_emit_cmd({"decision": "deny", "reason": "REPLACED OUTPUT"}),
        match="todo",
    )
    mistral_api.reply(
        mistral_completion("", tool_calls=TODO_TOOL_CALL), mistral_completion("done")
    )
    agent = build_e2e_agent_loop(
        config=build_test_vibe_config(enabled_tools=["todo"]),
        hook_config_result=HookConfigResult(hooks=[replace_output_hook], issues=[]),
    )

    [_ async for _ in agent.act("Show my todos")]

    assert "REPLACED OUTPUT" in mistral_api.model_facing_text(1)


@pytest.mark.asyncio
async def test_post_agent_turn_hook_injects_retry_user_message(
    mistral_api: MistralAPI,
) -> None:
    # A post-agent-turn hook that denies forever would loop, so this script
    # denies only on its first run, using a sentinel file to stay silent after.
    deny_decision = json.dumps({"decision": "deny", "reason": "please continue"})
    deny_once = textwrap.dedent(f"""
        import os, sys
        sentinel = "__post_turn_sentinel__"
        if not os.path.exists(sentinel):
            open(sentinel, "w").close()
            sys.stdout.write({deny_decision!r})
    """)
    retry_hook = HookConfig(
        name="retry",
        type=HookType.POST_AGENT_TURN,
        command=f"{sys.executable} -c {shlex.quote(deny_once)}",
    )
    mistral_api.reply(
        mistral_completion("first answer"), mistral_completion("second answer")
    )
    agent = build_e2e_agent_loop(
        config=build_test_vibe_config(enabled_tools=["todo"]),
        hook_config_result=HookConfigResult(hooks=[retry_hook], issues=[]),
    )

    events = [event async for event in agent.act("Hello")]

    assert "please continue" in mistral_api.model_facing_text(1)
    answers = [e.content for e in events if isinstance(e, AssistantEvent)]
    assert answers == ["first answer", "second answer"]