vibe/tests/test_agent_tool_call.py

from __future__ import annotations

import asyncio
import json
from typing import Any

import pytest

from tests.mock.utils import mock_llm_chunk
from tests.stubs.fake_backend import FakeBackend
from tests.stubs.fake_tool import FakeTool
from vibe.core.agent import Agent
from vibe.core.config import SessionLoggingConfig, VibeConfig
from vibe.core.modes import AgentMode
from vibe.core.tools.base import BaseToolConfig, ToolPermission
from vibe.core.tools.builtins.todo import TodoItem
from vibe.core.types import (
    ApprovalResponse,
    AssistantEvent,
    BaseEvent,
    FunctionCall,
    LLMMessage,
    Role,
    SyncApprovalCallback,
    ToolCall,
    ToolCallEvent,
    ToolResultEvent,
)


async def act_and_collect_events(agent: Agent, prompt: str) -> list[BaseEvent]:
    return [ev async for ev in agent.act(prompt)]


def make_config(todo_permission: ToolPermission = ToolPermission.ALWAYS) -> VibeConfig:
    return VibeConfig(
        session_logging=SessionLoggingConfig(enabled=False),
        auto_compact_threshold=0,
        enabled_tools=["todo"],
        tools={"todo": BaseToolConfig(permission=todo_permission)},
        system_prompt_id="tests",
        include_project_context=False,
        include_prompt_detail=False,
    )


def make_todo_tool_call(call_id: str, action: str = "read") -> ToolCall:
    return ToolCall(
        id=call_id,
        function=FunctionCall(name="todo", arguments=f'{{"action": "{action}"}}'),
    )


def make_agent(
    *,
    auto_approve: bool = True,
    todo_permission: ToolPermission = ToolPermission.ALWAYS,
    backend: FakeBackend,
    approval_callback: SyncApprovalCallback | None = None,
) -> Agent:
    mode = AgentMode.AUTO_APPROVE if auto_approve else AgentMode.DEFAULT
    agent = Agent(
        make_config(todo_permission=todo_permission), mode=mode, backend=backend
    )
    if approval_callback:
        agent.set_approval_callback(approval_callback)
    return agent


@pytest.mark.asyncio
async def test_single_tool_call_executes_under_auto_approve() -> None:
    mocked_tool_call_id = "call_1"
    tool_call = make_todo_tool_call(mocked_tool_call_id)
    backend = FakeBackend([
        mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call]),
        mock_llm_chunk(content="I retrieved 0 todos.", finish_reason="stop"),
    ])
    agent = make_agent(auto_approve=True, backend=backend)

    events = await act_and_collect_events(agent, "What's my todo list?")

    assert [type(e) for e in events] == [
        AssistantEvent,
        ToolCallEvent,
        ToolResultEvent,
        AssistantEvent,
    ]
    assert isinstance(events[0], AssistantEvent)
    assert events[0].content == "Let me check your todos."
    assert isinstance(events[1], ToolCallEvent)
    assert events[1].tool_name == "todo"
    assert isinstance(events[2], ToolResultEvent)
    assert events[2].error is None
    assert events[2].skipped is False
    assert events[2].result is not None
    assert isinstance(events[3], AssistantEvent)
    assert events[3].content == "I retrieved 0 todos."
    # check conversation history
    tool_msgs = [m for m in agent.messages if m.role == Role.tool]
    assert len(tool_msgs) == 1
    assert tool_msgs[-1].tool_call_id == mocked_tool_call_id
    assert "total_count" in (tool_msgs[-1].content or "")


@pytest.mark.asyncio
async def test_tool_call_requires_approval_if_not_auto_approved() -> None:
    agent = make_agent(
        auto_approve=False,
        todo_permission=ToolPermission.ASK,
        backend=FakeBackend([
            mock_llm_chunk(
                content="Let me check your todos.",
                tool_calls=[make_todo_tool_call("call_2")],
            ),
            mock_llm_chunk(
                content="I cannot execute the tool without approval.",
                finish_reason="stop",
            ),
        ]),
    )

    events = await act_and_collect_events(agent, "What's my todo list?")

    assert isinstance(events[1], ToolCallEvent)
    assert events[1].tool_name == "todo"
    assert isinstance(events[2], ToolResultEvent)
    assert events[2].skipped is True
    assert events[2].error is None
    assert events[2].result is None
    assert events[2].skip_reason is not None
    assert "not permitted" in events[2].skip_reason.lower()
    assert isinstance(events[3], AssistantEvent)
    assert events[3].content == "I cannot execute the tool without approval."
    assert agent.stats.tool_calls_rejected == 1
    assert agent.stats.tool_calls_agreed == 0
    assert agent.stats.tool_calls_succeeded == 0


@pytest.mark.asyncio
async def test_tool_call_approved_by_callback() -> None:
    def approval_callback(
        _tool_name: str, _args: dict[str, Any], _tool_call_id: str
    ) -> tuple[ApprovalResponse, str | None]:
        return (ApprovalResponse.YES, None)

    agent = make_agent(
        auto_approve=False,
        todo_permission=ToolPermission.ASK,
        approval_callback=approval_callback,
        backend=FakeBackend([
            mock_llm_chunk(
                content="Let me check your todos.",
                tool_calls=[make_todo_tool_call("call_3")],
            ),
            mock_llm_chunk(content="I retrieved 0 todos.", finish_reason="stop"),
        ]),
    )

    events = await act_and_collect_events(agent, "What's my todo list?")

    assert isinstance(events[2], ToolResultEvent)
    assert events[2].skipped is False
    assert events[2].error is None
    assert events[2].result is not None
    assert agent.stats.tool_calls_agreed == 1
    assert agent.stats.tool_calls_rejected == 0
    assert agent.stats.tool_calls_succeeded == 1


@pytest.mark.asyncio
async def test_tool_call_rejected_when_auto_approve_disabled_and_rejected_by_callback() -> (
    None
):
    custom_feedback = "User declined tool execution"

    def approval_callback(
        _tool_name: str, _args: dict[str, Any], _tool_call_id: str
    ) -> tuple[ApprovalResponse, str | None]:
        return (ApprovalResponse.NO, custom_feedback)

    agent = make_agent(
        auto_approve=False,
        todo_permission=ToolPermission.ASK,
        approval_callback=approval_callback,
        backend=FakeBackend([
            mock_llm_chunk(
                content="Let me check your todos.",
                tool_calls=[make_todo_tool_call("call_4")],
            ),
            mock_llm_chunk(
                content="Understood, I won't check the todos.", finish_reason="stop"
            ),
        ]),
    )

    events = await act_and_collect_events(agent, "What's my todo list?")

    assert isinstance(events[2], ToolResultEvent)
    assert events[2].skipped is True
    assert events[2].error is None
    assert events[2].result is None
    assert events[2].skip_reason == custom_feedback
    assert agent.stats.tool_calls_rejected == 1
    assert agent.stats.tool_calls_agreed == 0
    assert agent.stats.tool_calls_succeeded == 0


@pytest.mark.asyncio
async def test_tool_call_skipped_when_permission_is_never() -> None:
    agent = make_agent(
        auto_approve=False,
        todo_permission=ToolPermission.NEVER,
        backend=FakeBackend([
            mock_llm_chunk(
                content="Let me check your todos.",
                tool_calls=[make_todo_tool_call("call_never")],
            ),
            mock_llm_chunk(content="Tool is disabled.", finish_reason="stop"),
        ]),
    )

    events = await act_and_collect_events(agent, "What's my todo list?")

    assert isinstance(events[2], ToolResultEvent)
    assert events[2].skipped is True
    assert events[2].error is None
    assert events[2].result is None
    assert events[2].skip_reason is not None
    assert "permanently disabled" in events[2].skip_reason.lower()
    tool_msgs = [m for m in agent.messages if m.role == Role.tool and m.name == "todo"]
    assert len(tool_msgs) == 1
    assert tool_msgs[0].name == "todo"
    assert events[2].skip_reason in (tool_msgs[-1].content or "")
    assert agent.stats.tool_calls_rejected == 1
    assert agent.stats.tool_calls_agreed == 0
    assert agent.stats.tool_calls_succeeded == 0


@pytest.mark.asyncio
async def test_approval_always_sets_tool_permission_for_subsequent_calls() -> None:
    callback_invocations = []
    agent_ref: Agent | None = None

    def approval_callback(
        tool_name: str, _args: dict[str, Any], _tool_call_id: str
    ) -> tuple[ApprovalResponse, str | None]:
        callback_invocations.append(tool_name)
        # Set permission to ALWAYS for this tool (simulating the new behavior)
        assert agent_ref is not None
        if tool_name not in agent_ref.config.tools:
            agent_ref.config.tools[tool_name] = BaseToolConfig()
        agent_ref.config.tools[tool_name].permission = ToolPermission.ALWAYS
        return (ApprovalResponse.YES, None)

    agent = make_agent(
        auto_approve=False,
        todo_permission=ToolPermission.ASK,
        approval_callback=approval_callback,
        backend=FakeBackend([
            mock_llm_chunk(
                content="First check.", tool_calls=[make_todo_tool_call("call_first")]
            ),
            mock_llm_chunk(content="First done.", finish_reason="stop"),
            mock_llm_chunk(
                content="Second check.", tool_calls=[make_todo_tool_call("call_second")]
            ),
            mock_llm_chunk(content="Second done.", finish_reason="stop"),
        ]),
    )
    agent_ref = agent

    events1 = await act_and_collect_events(agent, "First request")
    events2 = await act_and_collect_events(agent, "Second request")

    tool_config_todo = agent.tool_manager.get_tool_config("todo")
    assert tool_config_todo.permission is ToolPermission.ALWAYS
    tool_config_help = agent.tool_manager.get_tool_config("bash")
    assert tool_config_help.permission is not ToolPermission.ALWAYS
    assert agent.auto_approve is False
    assert len(callback_invocations) == 1
    assert callback_invocations[0] == "todo"
    assert isinstance(events1[2], ToolResultEvent)
    assert events1[2].skipped is False
    assert events1[2].result is not None
    assert isinstance(events2[2], ToolResultEvent)
    assert events2[2].skipped is False
    assert events2[2].result is not None
    assert agent.stats.tool_calls_rejected == 0
    assert agent.stats.tool_calls_succeeded == 2


@pytest.mark.asyncio
async def test_tool_call_with_invalid_action() -> None:
    tool_call = ToolCall(
        id="call_5",
        function=FunctionCall(name="todo", arguments='{"action": "invalid_action"}'),
    )
    agent = make_agent(
        auto_approve=True,
        backend=FakeBackend([
            mock_llm_chunk(content="Let me check your todos.", tool_calls=[tool_call]),
            mock_llm_chunk(
                content="I encountered an error with the action.", finish_reason="stop"
            ),
        ]),
    )

    events = await act_and_collect_events(agent, "What's my todo list?")

    assert isinstance(events[2], ToolResultEvent)
    assert events[2].error is not None
    assert events[2].result is None
    assert "tool_error" in events[2].error.lower()
    assert agent.stats.tool_calls_failed == 1


@pytest.mark.asyncio
async def test_tool_call_with_duplicate_todo_ids() -> None:
    duplicate_todos = [
        TodoItem(id="duplicate", content="Task 1"),
        TodoItem(id="duplicate", content="Task 2"),
    ]
    tool_call = ToolCall(
        id="call_6",
        function=FunctionCall(
            name="todo",
            arguments=json.dumps({
                "action": "write",
                "todos": [t.model_dump() for t in duplicate_todos],
            }),
        ),
    )
    agent = make_agent(
        auto_approve=True,
        backend=FakeBackend([
            mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call]),
            mock_llm_chunk(
                content="I couldn't write todos with duplicate IDs.",
                finish_reason="stop",
            ),
        ]),
    )

    events = await act_and_collect_events(agent, "Add todos")

    assert isinstance(events[2], ToolResultEvent)
    assert events[2].error is not None
    assert events[2].result is None
    assert "unique" in events[2].error.lower()
    assert agent.stats.tool_calls_failed == 1


@pytest.mark.asyncio
async def test_tool_call_with_exceeding_max_todos() -> None:
    many_todos = [TodoItem(id=f"todo_{i}", content=f"Task {i}") for i in range(150)]
    tool_call = ToolCall(
        id="call_7",
        function=FunctionCall(
            name="todo",
            arguments=json.dumps({
                "action": "write",
                "todos": [t.model_dump() for t in many_todos],
            }),
        ),
    )
    agent = make_agent(
        auto_approve=True,
        backend=FakeBackend([
            mock_llm_chunk(content="Let me write todos.", tool_calls=[tool_call]),
            mock_llm_chunk(
                content="I couldn't write that many todos.", finish_reason="stop"
            ),
        ]),
    )

    events = await act_and_collect_events(agent, "Add todos")

    assert isinstance(events[2], ToolResultEvent)
    assert events[2].error is not None
    assert events[2].result is None
    assert "100" in events[2].error
    assert agent.stats.tool_calls_failed == 1


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "exception_class",
    [
        pytest.param(KeyboardInterrupt, id="keyboard_interrupt"),
        pytest.param(asyncio.CancelledError, id="asyncio_cancelled"),
    ],
)
async def test_tool_call_can_be_interrupted(
    exception_class: type[BaseException],
) -> None:
    tool_call = ToolCall(
        id="call_8", function=FunctionCall(name="stub_tool", arguments="{}")
    )
    config = VibeConfig(
        session_logging=SessionLoggingConfig(enabled=False),
        auto_compact_threshold=0,
        enabled_tools=["stub_tool"],
    )
    agent = Agent(
        config,
        mode=AgentMode.AUTO_APPROVE,
        backend=FakeBackend([
            mock_llm_chunk(content="Let me use the tool.", tool_calls=[tool_call]),
            mock_llm_chunk(content="Tool execution completed.", finish_reason="stop"),
        ]),
    )
    # no dependency injection available => monkey patch
    agent.tool_manager._available["stub_tool"] = FakeTool
    stub_tool_instance = agent.tool_manager.get("stub_tool")
    assert isinstance(stub_tool_instance, FakeTool)
    stub_tool_instance._exception_to_raise = exception_class()

    events: list[BaseEvent] = []
    with pytest.raises(exception_class):
        async for ev in agent.act("Execute tool"):
            events.append(ev)

    tool_result_event = next(
        (e for e in events if isinstance(e, ToolResultEvent)), None
    )
    assert tool_result_event is not None
    assert tool_result_event.error is not None
    assert "execution interrupted by user" in tool_result_event.error.lower()


@pytest.mark.asyncio
async def test_fill_missing_tool_responses_inserts_placeholders() -> None:
    agent = Agent(
        make_config(),
        mode=AgentMode.AUTO_APPROVE,
        backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
    )
    tool_calls_messages = [
        ToolCall(
            id="tc1", function=FunctionCall(name="todo", arguments='{"action": "read"}')
        ),
        ToolCall(
            id="tc2", function=FunctionCall(name="todo", arguments='{"action": "read"}')
        ),
    ]
    assistant_msg = LLMMessage(
        role=Role.assistant, content="Calling tools...", tool_calls=tool_calls_messages
    )
    agent.messages = [
        agent.messages[0],
        assistant_msg,
        # only one tool responded: the second is missing
        LLMMessage(
            role=Role.tool, tool_call_id="tc1", name="todo", content="Retrieved 0 todos"
        ),
    ]

    await act_and_collect_events(agent, "Proceed")

    tool_msgs = [m for m in agent.messages if m.role == Role.tool]
    assert any(m.tool_call_id == "tc2" for m in tool_msgs)
    # find placeholder message for tc2
    placeholder = next(m for m in tool_msgs if m.tool_call_id == "tc2")
    assert placeholder.name == "todo"
    assert (
        placeholder.content
        == "<user_cancellation>Tool execution interrupted - no response available</user_cancellation>"
    )


@pytest.mark.asyncio
async def test_ensure_assistant_after_tool_appends_understood() -> None:
    agent = Agent(
        make_config(),
        mode=AgentMode.AUTO_APPROVE,
        backend=FakeBackend([mock_llm_chunk(content="ok", finish_reason="stop")]),
    )
    tool_msg = LLMMessage(
        role=Role.tool, tool_call_id="tc_z", name="todo", content="Done"
    )
    agent.messages = [agent.messages[0], tool_msg]

    await act_and_collect_events(agent, "Next")

    # find the seeded tool message and ensure the next message is "Understood."
    idx = next(i for i, m in enumerate(agent.messages) if m.role == Role.tool)
    assert agent.messages[idx + 1].role == Role.assistant
    assert agent.messages[idx + 1].content == "Understood."