dev-puppeteer/my-deepagent/src/my_deepagent/middleware/plan_mode.py

"""PlanModeMiddleware (v0.3 PR #5) — block write tools when plan-mode is active.

Claude Code's plan mode lets the user say "design this, don't write code" — the
agent can read, search, plan via `write_todos`, but cannot mutate the
filesystem or run shell commands until the user `/approve`s.

Implementation strategy:
- A callable ``is_active()`` is passed in at construction time.  The REPL flips
  a flag on/off via slash commands; the middleware re-reads on every tool call.
  This avoids rebuilding the agent on every `/plan` / `/approve` toggle.
- When plan-mode is on and the LLM calls a blocked tool, we return a synthetic
  ``ToolMessage(status="error", ...)`` so the LLM sees feedback and can adjust
  ("ok, I'll keep planning instead").  We do NOT raise — that would crash the
  turn and the user would lose the partial response.

Blocked tools (matches Claude Code's ExitPlanMode-required tool set):
    - ``write_file``, ``edit_file`` — fs mutation
    - ``bash`` / ``execute`` / ``run_command`` / ``shell`` — shell exec
    - ``task`` — sub-agent spawn (a sub-agent could bypass plan mode)
    - ``write_todos`` — todos are PART of the plan markdown.  Plan-mode
      forbids commits to the agent's TODO list; the user reviews the plan
      first, then /approve unlocks both writes and the TODO list.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

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

#: Tool names that mutate the filesystem.
_FS_WRITE_TOOLS: frozenset[str] = frozenset({"write_file", "edit_file"})

#: Tool names that execute shell commands.
_SHELL_TOOLS: frozenset[str] = frozenset({"bash", "execute", "run_command", "shell"})

#: Tool names that spawn sub-agents (which would bypass plan mode in the parent).
_SUBAGENT_TOOLS: frozenset[str] = frozenset({"task"})

#: Plan-mode forbids committing to a TODO list — todos are part of the
#: plan markdown that the user reviews before /approve.
_PLANNING_TOOLS: frozenset[str] = frozenset({"write_todos"})

#: Full blocklist applied while plan mode is on.
BLOCKED_TOOLS_IN_PLAN_MODE: frozenset[str] = (
    _FS_WRITE_TOOLS | _SHELL_TOOLS | _SUBAGENT_TOOLS | _PLANNING_TOOLS
)


def _block_message(tool_name: str) -> str:
    return (
        f"Plan-mode is active — `{tool_name}` is blocked. "
        "Keep planning with read_file / glob / grep / write_todos, "
        "or ask the user to `/approve` to leave plan mode."
    )


class PlanModeMiddleware(AgentMiddleware):
    """Block mutating tool calls while plan-mode is active.

    Construction takes an ``is_active`` callable that returns the current plan
    mode state.  The REPL toggles this state via slash commands without
    rebuilding the agent — the middleware reads it fresh per tool call.

    Tools that are read-only (``read_file``, ``glob``, ``grep``, ``ls``,
    ``write_todos``) are allowed in plan mode unconditionally.
    """

    def __init__(self, *, is_active: Callable[[], bool]) -> None:
        self._is_active = is_active

    async def awrap_tool_call(self, request: Any, handler: Any) -> Any:
        if not self._is_active():
            return await handler(request)
        name = _tool_name(request)
        if name in BLOCKED_TOOLS_IN_PLAN_MODE:
            return ToolMessage(
                content=_block_message(name),
                tool_call_id=_tool_call_id(request),
                name=name,
                status="error",
            )
        return await handler(request)

    def wrap_tool_call(self, request: Any, handler: Any) -> Any:
        # Sync path mirrors the async one for parity (e.g. when the agent is
        # invoked synchronously in unit tests).  Real REPL/Web paths are async.
        if not self._is_active():
            return handler(request)
        name = _tool_name(request)
        if name in BLOCKED_TOOLS_IN_PLAN_MODE:
            return ToolMessage(
                content=_block_message(name),
                tool_call_id=_tool_call_id(request),
                name=name,
                status="error",
            )
        return handler(request)


def _tool_name(request: Any) -> str:
    tool_call = getattr(request, "tool_call", None)
    if isinstance(tool_call, dict):
        return str(tool_call.get("name") or "")
    return str(getattr(request, "name", "") or "")


def _tool_call_id(request: Any) -> str:
    tool_call = getattr(request, "tool_call", None)
    if isinstance(tool_call, dict):
        return str(tool_call.get("id") or "")
    return str(getattr(request, "id", "") or "")