Planning & TODOs: A Working Memory for Long Tasks

Ask Axon for something that takes more than a couple of steps and it starts well, then quietly wanders. It reads a file, gets interested in what it found, and finishes something next to what you asked for. Your request is still in the conversation. It is just far away by then, with thirty tool calls piled on top of it.

This chapter we add the feature that lets the agent plan for itself: a TODO list it writes, re-reads, and rewrites as it works. Two tools, one field on State.

1. One request, start to finish#

This project still has no README, so write one.

Press play, or step with the arrows. The graph on top shows where the turn is at every step. human_approval runs in auto mode here, so calls pass straight through it (that switch is section 8). Below it, the context window is the actual message list the model sees, growing one message at a time: each AIMessage carries the calls the model produced, and each ToolMessage carries what came back. The labels above the progress bar mark the phases of the run, and the tall ticks are the key steps. Click either to jump.

You send one message. LLM answers with either a tool call or a final answer. A tool returns text that goes straight back into the thread.

You are Axon, a terminal-native coding assistant.

Two things are worth taking from that run before we open any code.

The thread is the only thing that carries the plan. Watch where the write_todos output lands: at the bottom, every time, which is exactly where the model looks next. There is no plan variable, no queue, no scheduler. By the last step the thread is 19 messages and about 2,500 characters, and the plan has been posted into it four times. The plan is N messages from the end counter in the state bar makes the rhythm visible: it climbs while the agent reads files, and snaps back to zero on every rewrite. That climb-and-reset is recitation, measured.

Almost none of it was work. Nine model turns, and exactly one of them changed a file. Four of them rewrote a list that no code anywhere reads. That looks wasteful right up until you ask what the list is actually for, which is section 3.

2. The graph#

Same nodes, same edges as chapter 3; go back to the graph there if you want to click through it. Planning did not add a planner node, and did not add an edge.

What did change is the contract each node has with state. Only tools touches todos; chat, the node that actually talks to the model, does not read it at all.

So the plan has no way to reach the model by sitting in state. It reaches the model only as text that is already in the message list, which is exactly what the run above looks like. Everything in the rest of this chapter follows from that one fact.

Four pieces get us there:

3. The write_todos tool#

A task is two fields. Nothing more:

class Todo(TypedDict):
    content: str
    status: Literal["pending", "in_progress", "completed"]

No ids, no dependencies, no priorities. This is a note to self, not a project tracker.

Here is the whole tool:

STATUS_ICON = {"pending": "[ ]", "in_progress": "[~]", "completed": "[x]"}
 
 
def render_todos(todos: list[Todo]) -> str:
    return "\n".join(
        f"{STATUS_ICON.get(t['status'], '[?]')} {t['content']}" for t in todos
    )
 
 
@tool
def write_todos(todos: list[Todo]) -> str:
    """Create or replace the task list used to plan multi-step work.
 
    Use this at the start of any non-trivial request to break it into concrete
    steps, then call it again after every step to update progress. The list is
    your working memory for long tasks -- rewriting it keeps you from drifting
    off the original goal over a long session.
 
    Do NOT use this for single, trivial actions (reading one file, answering a
    question). The overhead is not worth it.
 
    Rules:
        - Send the FULL list every time; this replaces the previous list.
          Revise, merge, or drop steps freely as you learn more.
        - Keep exactly one step 'in_progress' at a time.
        - Mark a step 'completed' the moment it is done, not in batches.
        - If a step is blocked, leave it 'in_progress' and add a new step
          describing what unblocks it.
 
    Args:
        todos: The complete task list. Each item needs 'content' (a short,
            actionable description) and 'status'.
 
    Returns:
        The rendered task list as it was saved.
    """
    if not todos:
        return "Task list cleared."
    return "Task list updated:\n" + render_todos(todos)

Read the body again: it formats a string and returns it. It runs nothing, changes nothing, enforces nothing. The tool is a no-op.

That is not a simplification for the article, it is the actual design, and the run in section 1 shows why it works anyway. Go back to any write_todos step in section 1 and step forward once. The tool returned entry is a copy of the call the llm made a moment earlier, and it lands at the bottom of the thread. That is the entire effect of the call. Two things happened there, and neither of them is mechanical:

1. Decomposition. To call the tool at all, the model had to break the request into concrete steps before starting. That alone improves the result.
2. Recitation. The returned string enters the conversation as a ToolMessage, so after every update the full plan reappears at the end of the context, which is the position the model attends to most. The agent is reminding itself out loud.

Because the tool has no mechanical power, all of the engineering sits in the docstring, which is the only thing the model ever reads about it. Three lines there do real work:

• "Send the FULL list every time" is why the parameter is the whole list instead of a patch. Rewriting invites the model to reconsider the plan as it learns, rather than mechanically ticking off a stale one.
• "Keep exactly one step in_progress" keeps the list readable and forces the model to commit to what it is doing right now.
• "Do NOT use this for single, trivial actions" is the off switch. Without it the agent writes a three item plan to answer "what does this function do?".

4. The read_todos tool#

Recitation puts the plan in context at the moment it is written, but twenty tool calls later it has scrolled away again. So the agent needs to be able to pull it back:

@tool
def read_todos(state: Annotated[dict, InjectedState]) -> str:
    """Re-read the current task list to check what is done and what is left.
 
    Call this after finishing a step, before deciding what to do next. On long
    tasks the plan scrolls far back in the conversation; this brings it back
    into view so you stay on the original goal.
 
    Returns:
        The task list with a status marker per step, or a note that no list
        exists yet (in which case call write_todos first).
    """
    todos = state.get("todos") or []
    if not todos:
        return "No task list yet. Call write_todos to create one."
    return render_todos(todos)

The two tools get their data from opposite directions, and that is the interesting part. write_todos is handed its list by the model, which literally types it out. read_todos needs the list from the agent's state, which the model has never seen and must not be asked to supply.

That is the read_todos step in the run above: the AIMessage carries read_todos() and nothing else, and the ToolMessage under it is tagged with the state that tool_node injected on the way in. Annotated[dict, InjectedState] is what makes that split legal.

5. InjectedState: how a tool reads agent state#

Every tool has two consumers, and it is easy to miss that they are different:

1. The LLM reads the tool's JSON schema and fills in the arguments.
2. The runtime actually calls the Python function.

Usually they see the same thing. For read_file(path) the model supplies path, the runtime calls with path, done.

read_todos breaks that. If we write def read_todos(state) with no annotation, the @tool decorator builds the schema from the type hints, the model sees a state parameter, and it tries to invent the entire agent state as an argument.

InjectedState remove the parameter from the schema sent to the model, keep it in the tool signature.

Below are those two views for three tools. Switch between them and watch the footer count. For the first two the views agree, which is why no toggle appears. For read_todos they come apart, and the toggle lets you take the annotation away and watch state leak into the model's schema.

{
"path": "string",
"offset": "integer (optional)",
"limit": "integer (optional)"
}

def read_file(
  path: str,
  offset: int = 1,
  limit: int = 2000,
) -> str:

Zero parameters on the left, one on the right. Something in the middle has to fill that gap, which is section 7.

Three properties are worth knowing before you use it:

It does not reach the model. state is a Python function argument, not context. The only thing that goes back to the LLM is the tool's return value. Injecting state costs no tokens and does not resend the conversation.

It passes the whole state, by reference. Not a copy, not a filtered view: the tool can see every channel, and a careless state["todos"].append(...) would mutate the real state behind LangGraph's back, without going through a reducer or leaving a trace in the checkpoint history. Treat it as read-only.

It does not create the field. It passes along whatever happens to be in state at that moment. Declaring the field is a separate job, and that is the next section.

6. Add the todos channel#

Our state has been messages and nothing else since chapter 1:

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

Run the todo tools against that and nothing crashes, which is exactly the problem. Here is a node returning a key the state does not declare:

def writes_todos(state):
    return {"messages": [], "todos": TODOS}

result: {}
did todos land in state? -> False

No exception, no warning. LangGraph silently drops it. State is not a free form dict, it is a fixed set of declared channels, and an update to a channel that does not exist goes nowhere. The read side gives you the mirror image: read_todos returns "No task list yet" forever, and the agent happily writes the plan again.

One line fixes it:

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    todos: list[Todo]

Note there is no reducer on todos, unlike messages with its add_messages. That is deliberate, and it is the same decision as "send the full list every time": a write replaces the list wholesale, so the model can reorganize the plan instead of only appending to it.

7. Wiring it into tool_node#

Two things have to happen in the tool loop now: injected parameters need filling, and a write_todos call needs to land in state. Here is the finished node, with the new lines marked:

def tool_node(state: State) -> dict:
    ai = next(m for m in reversed(state["messages"]) if getattr(m, "tool_calls", None))
    answered_ids = {m.tool_call_id for m in state["messages"] if isinstance(m, ToolMessage)}
    working = state                                  # NEW: advances during the turn
    update = {}                                      # NEW: non-message state updates
    messages = []
    for tc in ai.tool_calls:
        if tc["id"] in answered_ids:                 # rejected in human_approval -> skip
            continue
        tool = TOOLS_BY_NAME[tc["name"]]
        result = tool.invoke({**tc["args"], "state": working})        # NEW: injection
        if tc["name"] == "write_todos":                               # NEW: the write
            update["todos"] = tc["args"]["todos"]
            working = {**working, "todos": tc["args"]["todos"]}
        messages.append(ToolMessage(str(result), tool_call_id=tc["id"], name=tc["name"]))
    return {"messages": messages, **update}

Injection. In a prebuilt agent, ToolNode fills injected parameters for you. We wrote our own node back in chapter 2, so the job is ours, and it is one expression: {**tc["args"], "state": working}. That expression is the "added by tool_node" line you saw in the read_todos step. We pass state to every tool rather than checking which ones asked, because Pydantic silently drops unexpected keys. If you prefer to be strict about it:

def needs_state(tool) -> bool:
    return "state" in tool.get_input_schema().model_fields

get_input_schema() is the real signature and tool.args is the LLM-facing one. Those are the two panes of the lens in section 5, and where they differ is precisely the set of injected parameters.

The write. write_todos returns a plain string, so the node lifts the list straight out of the call arguments. No Command and no InjectedToolCallId are needed; those exist for tools that must reach state from outside the tool loop, and we are inside it, where state and tc["id"] are ordinary local variables.

Two accumulators. messages collects ToolMessages as before, update collects everything else. They merge on the way out with {"messages": messages, **update}.

That leaves working, which looks redundant next to update but guards against a real failure. Models emit parallel tool calls, so one AIMessage can carry write_todos and read_todos together, and core LangGraph semantics say:

The state handed to a node is a snapshot. The update a node returns is applied after the node finishes.

So inside one tool_node execution, state["todos"] stays frozen: a write_todos two lines earlier has not landed yet, and a naive read_todos in the same turn would answer "No task list yet." working starts as the snapshot and advances as calls execute, so a read later in the turn sees a write from earlier in it, while update carries the real change back to LangGraph at the end. Two jobs, two variables.

8. Auto mode#

Chapter 3 asks you to approve every tool call. With todo tools in the mix that gets absurd, since you end up approving a function that does nothing. Section 1 would have needed eight extra confirmations.

The tempting fix is a SAFE_TOOLS set that skips approval for anything read-only. We did not do that, because it means the author of the agent decides what counts as safe on your behalf, silently, in a set literal. Approval stays all or nothing, and you turn it off yourself:

def human_approval_node(state: State, config: RunnableConfig) -> dict:
    if config["configurable"].get("auto_approve"):      # /auto -> run without asking
        return {"messages": []}
    ...                                                 # unchanged from chapter 3

The flag lives in config, not State, because it is a runtime setting rather than something the agent knows about itself. The CLI already builds a config every turn, and it should not be checkpointed alongside the conversation.

if text == "/auto":
    auto = not config["configurable"]["auto_approve"]
    config["configurable"]["auto_approve"] = auto
    print(f"auto-approve {'ON' if auto else 'OFF'}")
    continue

9. Up next#

The agent can hold a plan while it works. What it cannot do yet is keep its own context clean while it works: every file it reads to check one fact stays in the thread for the rest of the session.

Next chapter: subagents. The same loop, started on a fresh message list with its own tools and prompt, doing the reading somewhere else and sending back a report. Plan mode is still owed; it turns out to be a configuration of what we build next.