dev-puppeteer/my-deepagent/tests/integration/test_plan_mode.py

"""v0.3 PR #5 — Plan mode tests.

Covers:
1. PlanModeMiddleware passes tool calls through when inactive.
2. PlanModeMiddleware blocks write_file / edit_file / execute / task when active.
3. read_file / glob / grep / write_todos are allowed regardless.
4. Toggling the closure flag changes behavior without rebuilding the middleware.
5. The synthetic ToolMessage carries status="error" and a clear hint.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import pytest
from langchain_core.messages import ToolMessage

from my_deepagent.middleware.plan_mode import (
    BLOCKED_TOOLS_IN_PLAN_MODE,
    PlanModeMiddleware,
)


@dataclass
class _FakeToolRequest:
    """Minimal stand-in for langchain ToolCallRequest in unit tests."""

    tool_call: dict[str, Any]


async def _passthrough_handler(_: _FakeToolRequest) -> ToolMessage:
    """Stub handler — returns a benign 'tool executed' message."""
    return ToolMessage(content="EXECUTED", tool_call_id="t1", name="stub")


# ---------------------------------------------------------------------------
# Inactive plan-mode → all tools pass through
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plan_mode_inactive_passes_through() -> None:
    mw = PlanModeMiddleware(is_active=lambda: False)
    for name in ["write_file", "edit_file", "execute", "task", "read_file", "glob"]:
        req = _FakeToolRequest(tool_call={"name": name, "id": "t1", "args": {}})
        result = await mw.awrap_tool_call(req, _passthrough_handler)
        assert isinstance(result, ToolMessage)
        assert result.content == "EXECUTED"
        assert result.status != "error"


# ---------------------------------------------------------------------------
# Active plan-mode → write tools blocked with status=error
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plan_mode_active_blocks_write_file() -> None:
    mw = PlanModeMiddleware(is_active=lambda: True)
    req = _FakeToolRequest(
        tool_call={"name": "write_file", "id": "abc123", "args": {"file_path": "/tmp/x"}}
    )
    result = await mw.awrap_tool_call(req, _passthrough_handler)
    assert isinstance(result, ToolMessage)
    assert result.status == "error"
    assert result.tool_call_id == "abc123"
    assert "Plan-mode" in result.content
    assert "write_file" in result.content


@pytest.mark.asyncio
async def test_plan_mode_active_blocks_execute() -> None:
    mw = PlanModeMiddleware(is_active=lambda: True)
    req = _FakeToolRequest(tool_call={"name": "execute", "id": "exec1", "args": {"command": "ls"}})
    result = await mw.awrap_tool_call(req, _passthrough_handler)
    assert isinstance(result, ToolMessage)
    assert result.status == "error"
    assert "execute" in result.content


@pytest.mark.asyncio
async def test_plan_mode_active_blocks_task_subagent_spawn() -> None:
    mw = PlanModeMiddleware(is_active=lambda: True)
    req = _FakeToolRequest(tool_call={"name": "task", "id": "task1", "args": {"description": "x"}})
    result = await mw.awrap_tool_call(req, _passthrough_handler)
    assert isinstance(result, ToolMessage)
    assert result.status == "error"
    assert "task" in result.content


# ---------------------------------------------------------------------------
# Active plan-mode → read-only tools still pass through
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plan_mode_active_allows_read_only_tools() -> None:
    mw = PlanModeMiddleware(is_active=lambda: True)
    for name in ["read_file", "glob", "grep", "ls"]:
        req = _FakeToolRequest(tool_call={"name": name, "id": "t1", "args": {}})
        result = await mw.awrap_tool_call(req, _passthrough_handler)
        assert result.content == "EXECUTED", f"{name} should not be blocked"
        assert result.status != "error"


@pytest.mark.asyncio
async def test_plan_mode_blocks_write_todos() -> None:
    """`write_todos` is part of the plan markdown — must be blocked."""
    mw = PlanModeMiddleware(is_active=lambda: True)
    req = _FakeToolRequest(tool_call={"name": "write_todos", "id": "wt1", "args": {"todos": []}})
    result = await mw.awrap_tool_call(req, _passthrough_handler)
    assert isinstance(result, ToolMessage)
    assert result.status == "error"
    assert "write_todos" in result.content


# ---------------------------------------------------------------------------
# Closure-toggle behavior — flip without rebuild
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_plan_mode_closure_toggle_changes_behavior() -> None:
    state = {"on": False}
    mw = PlanModeMiddleware(is_active=lambda: state["on"])

    req = _FakeToolRequest(tool_call={"name": "write_file", "id": "w", "args": {}})

    # Off → passes.
    r1 = await mw.awrap_tool_call(req, _passthrough_handler)
    assert r1.status != "error"

    # Flip on → blocks.
    state["on"] = True
    r2 = await mw.awrap_tool_call(req, _passthrough_handler)
    assert r2.status == "error"

    # Flip back off → passes again.
    state["on"] = False
    r3 = await mw.awrap_tool_call(req, _passthrough_handler)
    assert r3.status != "error"


# ---------------------------------------------------------------------------
# Sync path mirrors async path
# ---------------------------------------------------------------------------


def test_plan_mode_sync_wrap_tool_call() -> None:
    mw = PlanModeMiddleware(is_active=lambda: True)

    def sync_handler(_: _FakeToolRequest) -> ToolMessage:
        return ToolMessage(content="EXECUTED", tool_call_id="t1", name="stub")

    req = _FakeToolRequest(tool_call={"name": "write_file", "id": "s1", "args": {}})
    result = mw.wrap_tool_call(req, sync_handler)
    assert isinstance(result, ToolMessage)
    assert result.status == "error"


# ---------------------------------------------------------------------------
# Blocklist constant sanity
# ---------------------------------------------------------------------------


def test_blocklist_includes_all_known_write_tools() -> None:
    assert "write_file" in BLOCKED_TOOLS_IN_PLAN_MODE
    assert "edit_file" in BLOCKED_TOOLS_IN_PLAN_MODE
    assert "execute" in BLOCKED_TOOLS_IN_PLAN_MODE
    assert "bash" in BLOCKED_TOOLS_IN_PLAN_MODE
    assert "task" in BLOCKED_TOOLS_IN_PLAN_MODE


def test_blocklist_excludes_read_only_tools() -> None:
    for name in ("read_file", "glob", "grep", "ls"):
        assert name not in BLOCKED_TOOLS_IN_PLAN_MODE


def test_blocklist_includes_write_todos() -> None:
    assert "write_todos" in BLOCKED_TOOLS_IN_PLAN_MODE