1"""Planning and task management middleware for agents."""23from collections.abc import Awaitable, Callable4from typing import Annotated, Any, Literal, cast56from langchain_core.messages import AIMessage, SystemMessage, ToolMessage7from langchain_core.tools import InjectedToolCallId, StructuredTool, tool8from langgraph.runtime import Runtime9from langgraph.types import Command10from pydantic import BaseModel11from typing_extensions import NotRequired, TypedDict, override1213from langchain.agents.middleware.types import (14    AgentMiddleware,15    AgentState,16    ContextT,17    ModelRequest,18    ModelResponse,19    OmitFromInput,20    ResponseT,21)22from langchain.tools import ToolRuntime232425class Todo(TypedDict):26    """A single todo item with content and status."""2728    content: str29    """The content/description of the todo item."""3031    status: Literal["pending", "in_progress", "completed"]32    """The current status of the todo item."""333435class PlanningState(AgentState[ResponseT]):36    """State schema for the todo middleware.3738    Type Parameters:39        ResponseT: The type of the structured response. Defaults to `Any`.40    """4142    todos: Annotated[NotRequired[list[Todo]], OmitFromInput]43    """List of todo items for tracking task progress."""444546class WriteTodosInput(BaseModel):47    """Input schema for the `write_todos` tool."""4849    todos: list[Todo]505152WRITE_TODOS_TOOL_DESCRIPTION = """Use this tool to create and manage a structured task list for your current work session. This helps you track progress and organize complex tasks.5354Only use this tool if you think it will be helpful in staying organized. If the user's request is trivial and takes less than 3 steps, it is better to NOT use this tool and just do the task directly.5556## When to Use This Tool5758Use this tool in these scenarios:59601. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions612. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations623. User explicitly requests todo list - When the user directly asks you to use the todo list634. User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)645. The plan may need future revisions or updates based on results from the first few steps6566## How to Use This Tool67681. When you start working on a task - Mark it as in_progress BEFORE beginning work.692. After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation.703. You can also update future tasks, such as deleting them if they are no longer necessary, or adding new tasks that are necessary. Don't change previously completed tasks.714. You can make several updates to the todo list at once. For example, when you complete a task, you can mark the next task you need to start as in_progress.7273## When NOT to Use This Tool7475It is important to skip using this tool when:761. There is only a single, straightforward task772. The task is trivial and tracking it provides no benefit783. The task can be completed in less than 3 trivial steps794. The task is purely conversational or informational8081## Task States and Management82831. **Task States**: Use these states to track progress:84    - pending: Task not yet started85    - in_progress: Currently working on (you can have multiple tasks in_progress at a time if they are not related to each other and can be run in parallel)86    - completed: Task finished successfully87882. **Task Management**:89    - Update task status in real-time as you work90    - Mark tasks complete IMMEDIATELY after finishing (don't batch completions)91    - Complete current tasks before starting new ones92    - Remove tasks that are no longer relevant from the list entirely93    - IMPORTANT: When you write this todo list, you should mark your first task (or tasks) as in_progress immediately!.94    - IMPORTANT: Unless all tasks are completed, you should always have at least one task in_progress.95963. **Task Completion Requirements**:97    - ONLY mark a task as completed when you have FULLY accomplished it98    - If you encounter errors, blockers, or cannot finish, keep the task as in_progress99    - When blocked, create a new task describing what needs to be resolved100    - Never mark a task as completed if:101        - There are unresolved issues or errors102        - Work is partial or incomplete103        - You encountered blockers that prevent completion104        - You couldn't find necessary resources or dependencies105        - Quality standards haven't been met1061074. **Task Breakdown**:108    - Create specific, actionable items109    - Break complex tasks into smaller, manageable steps110    - Use clear, descriptive task names111112Being proactive with task management ensures you complete all requirements successfully113Remember: If you only need to make a few tool calls to complete a task, and it is clear what you need to do, it is better to just do the task directly and NOT call this tool at all.114115## When You Finish116117`write_todos` tracks your work; it does not deliver the answer. Whatever the user asked for — computations, summaries, comparisons, data — must appear as text content in a message after your final `write_todos` call. Marking the last todo complete is not itself an answer to the user."""  # noqa: E501118119WRITE_TODOS_SYSTEM_PROMPT = """## `write_todos`120121You have access to the `write_todos` tool to help you manage and plan complex objectives.122Use this tool for complex objectives to ensure that you are tracking each necessary step.123This tool is very helpful for planning complex objectives, and for breaking down these larger complex objectives into smaller steps.124125It is critical that you mark todos as completed as soon as you are done with a step. Do not batch up multiple steps before marking them as completed.126For simple objectives that only require a few steps, it is better to just complete the objective directly and NOT use this tool.127Writing todos takes time and tokens, use it when it is helpful for managing complex many-step problems! But not for simple few-step requests.128129## Important To-Do List Usage Notes to Remember130131- The `write_todos` tool should never be called multiple times in parallel.132- Don't be afraid to revise the To-Do list as you go. New information may reveal new tasks that need to be done, or old tasks that are irrelevant.133134## Finishing a task135136When you finish all work, write your final answer in the message AFTER your last `write_todos` call — not in the same turn as that call. Start the final message with the substantive content the user asked for — the data, computation, summary, or analysis. The user wants the result, not confirmation that the work is done."""  # noqa: E501137138139@tool(description=WRITE_TODOS_TOOL_DESCRIPTION)140def write_todos(141    todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]142) -> Command[Any]:143    """Create and manage a structured task list for your current work session."""144    return Command(145        update={146            "todos": todos,147            "messages": [ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id)],148        }149    )150151152# Dynamically create the write_todos tool with the custom description153def _write_todos(154    runtime: ToolRuntime[ContextT, PlanningState[ResponseT]], todos: list[Todo]155) -> Command[Any]:156    """Create and manage a structured task list for your current work session."""157    return Command(158        update={159            "todos": todos,160            "messages": [161                ToolMessage(f"Updated todo list to {todos}", tool_call_id=runtime.tool_call_id)162            ],163        }164    )165166167async def _awrite_todos(168    runtime: ToolRuntime[ContextT, PlanningState[ResponseT]], todos: list[Todo]169) -> Command[Any]:170    """Create and manage a structured task list for your current work session."""171    return _write_todos(runtime, todos)172173174class TodoListMiddleware(AgentMiddleware[PlanningState[ResponseT], ContextT, ResponseT]):175    """Middleware that provides todo list management capabilities to agents.176177    This middleware adds a `write_todos` tool that allows agents to create and manage178    structured task lists for complex multi-step operations. It's designed to help179    agents track progress, organize complex tasks, and provide users with visibility180    into task completion status.181182    The middleware automatically injects system prompts that guide the agent on when183    and how to use the todo functionality effectively. It also enforces that the184    `write_todos` tool is called at most once per model turn, since the tool replaces185    the entire todo list and parallel calls would create ambiguity about precedence.186187    Example:188        ```python189        from langchain.agents.middleware import TodoListMiddleware190        from langchain.agents import create_agent191192        agent = create_agent("openai:gpt-5.5", middleware=[TodoListMiddleware()])193194        # Agent now has access to write_todos tool and todo state tracking195        result = await agent.invoke({"messages": [HumanMessage("Help me refactor my codebase")]})196197        print(result["todos"])  # Array of todo items with status tracking198        ```199    """200201    state_schema = PlanningState  # type: ignore[assignment]202203    def __init__(204        self,205        *,206        system_prompt: str = WRITE_TODOS_SYSTEM_PROMPT,207        tool_description: str = WRITE_TODOS_TOOL_DESCRIPTION,208    ) -> None:209        """Initialize the `TodoListMiddleware` with optional custom prompts.210211        Args:212            system_prompt: Custom system prompt to guide the agent on using the todo213                tool.214            tool_description: Custom description for the `write_todos` tool.215        """216        super().__init__()217        self.system_prompt = system_prompt218        self.tool_description = tool_description219220        self.tools = [221            StructuredTool.from_function(222                name="write_todos",223                description=tool_description,224                func=_write_todos,225                coroutine=_awrite_todos,226                args_schema=WriteTodosInput,227                infer_schema=False,228            )229        ]230231    def wrap_model_call(232        self,233        request: ModelRequest[ContextT],234        handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],235    ) -> ModelResponse[ResponseT] | AIMessage:236        """Update the system message to include the todo system prompt.237238        Args:239            request: Model request to execute (includes state and runtime).240            handler: Async callback that executes the model request and returns241                `ModelResponse`.242243        Returns:244            The model call result.245        """246        if request.system_message is not None:247            new_system_content = [248                *request.system_message.content_blocks,249                {"type": "text", "text": f"\n\n{self.system_prompt}"},250            ]251        else:252            new_system_content = [{"type": "text", "text": self.system_prompt}]253        new_system_message = SystemMessage(254            content=cast("list[str | dict[str, str]]", new_system_content)255        )256        return handler(request.override(system_message=new_system_message))257258    async def awrap_model_call(259        self,260        request: ModelRequest[ContextT],261        handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],262    ) -> ModelResponse[ResponseT] | AIMessage:263        """Update the system message to include the todo system prompt.264265        Args:266            request: Model request to execute (includes state and runtime).267            handler: Async callback that executes the model request and returns268                `ModelResponse`.269270        Returns:271            The model call result.272        """273        if request.system_message is not None:274            new_system_content = [275                *request.system_message.content_blocks,276                {"type": "text", "text": f"\n\n{self.system_prompt}"},277            ]278        else:279            new_system_content = [{"type": "text", "text": self.system_prompt}]280        new_system_message = SystemMessage(281            content=cast("list[str | dict[str, str]]", new_system_content)282        )283        return await handler(request.override(system_message=new_system_message))284285    @override286    def after_model(287        self, state: PlanningState[ResponseT], runtime: Runtime[ContextT]288    ) -> dict[str, Any] | None:289        """Check for parallel write_todos tool calls and return errors if detected.290291        The todo list is designed to be updated at most once per model turn. Since292        the `write_todos` tool replaces the entire todo list with each call, making293        multiple parallel calls would create ambiguity about which update should take294        precedence. This method prevents such conflicts by rejecting any response that295        contains multiple write_todos tool calls.296297        Args:298            state: The current agent state containing messages.299            runtime: The LangGraph runtime instance.300301        Returns:302            A dict containing error ToolMessages for each write_todos call if multiple303            parallel calls are detected, otherwise None to allow normal execution.304        """305        messages = state["messages"]306        if not messages:307            return None308309        last_ai_msg = next((msg for msg in reversed(messages) if isinstance(msg, AIMessage)), None)310        if not last_ai_msg or not last_ai_msg.tool_calls:311            return None312313        # Count write_todos tool calls314        write_todos_calls = [tc for tc in last_ai_msg.tool_calls if tc["name"] == "write_todos"]315316        if len(write_todos_calls) > 1:317            # Create error tool messages for all write_todos calls318            error_messages = [319                ToolMessage(320                    content=(321                        "Error: The `write_todos` tool should never be called multiple times "322                        "in parallel. Please call it only once per model invocation to update "323                        "the todo list."324                    ),325                    tool_call_id=tc["id"],326                    status="error",327                )328                for tc in write_todos_calls329            ]330331            # Keep the tool calls in the AI message but return error messages332            # This follows the same pattern as HumanInTheLoopMiddleware333            return {"messages": error_messages}334335        return None336337    @override338    async def aafter_model(339        self, state: PlanningState[ResponseT], runtime: Runtime[ContextT]340    ) -> dict[str, Any] | None:341        """Check for parallel write_todos tool calls and return errors if detected.342343        Async version of `after_model`. The todo list is designed to be updated at344        most once per model turn. Since the `write_todos` tool replaces the entire345        todo list with each call, making multiple parallel calls would create ambiguity346        about which update should take precedence. This method prevents such conflicts347        by rejecting any response that contains multiple write_todos tool calls.348349        Args:350            state: The current agent state containing messages.351            runtime: The LangGraph runtime instance.352353        Returns:354            A dict containing error ToolMessages for each write_todos call if multiple355            parallel calls are detected, otherwise None to allow normal execution.356        """357        return self.after_model(state, runtime)