Working Memory: Scratchpads, Blackboards, and Agent Notebooks

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

# pip install "langgraph>=0.2" "langchain-anthropic>=0.2" "anthropic>=0.40"
import operator
from typing import Annotated, TypedDict
from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
from langchain_anthropic import ChatAnthropic
from langgraph.graph import StateGraph, START, END

# Two channels with different reducers — the dataflow-graph idea made concrete.
# `messages` accumulates (operator.add appends list-to-list).
# `plan`, `completed`, `notes` overwrite on each write; this is the working-memory tier
# the harness keeps explicit, separate from the chat buffer.
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]
    plan: list[str]              # overwrite — the model rewrites the plan as it learns
    completed: list[str]         # overwrite — the model emits the full list each time
    notes: Annotated[list[str], operator.add]  # append — facts pinned across turns

llm = ChatAnthropic(model="claude-opus-4-7", max_tokens=2048)

def working_memory_prompt(state: AgentState) -> SystemMessage:
    """Render the working-memory block that gets injected on every model call.
    This is the bit that distinguishes working memory from short-term memory: it's
    a deliberate, structured slab the harness assembles, not whatever the model
    happened to say in the last 20 turns."""
    lines = ["## Working memory (your private scratchpad — re-injected every turn)"]
    if state.get("plan"):
        lines.append("### Current plan")
        lines.extend(f"- {step}" for step in state["plan"])
    if state.get("completed"):
        lines.append("### Completed steps")
        lines.extend(f"- {step}" for step in state["completed"])
    if state.get("notes"):
        lines.append("### Pinned notes")
        lines.extend(f"- {note}" for note in state["notes"])
    return SystemMessage(content="\n".join(lines))

def model_node(state: AgentState) -> dict:
    """Standard LLM step. The working-memory block is prepended to the chat
    history; the model can update working memory by emitting a structured
    directive in its next tool call (omitted here for brevity)."""
    wm_block = working_memory_prompt(state)
    response = llm.invoke([wm_block] + state["messages"])
    return {"messages": [response]}

graph = StateGraph(AgentState)
graph.add_node("model", model_node)
graph.add_edge(START, "model")
graph.add_edge("model", END)
compiled = graph.compile()

# Initial state — the working-memory fields are explicit, separate from `messages`.
result = compiled.invoke({
    "messages": [HumanMessage("Refactor utils.py to remove the global state.")],
    "plan": [
        "Identify every reader of the global.",
        "Introduce a Config dataclass.",
        "Thread Config through call sites.",
        "Delete the global.",
    ],
    "completed": [],
    "notes": ["The global is named CONFIG and is set in utils.py:14."],
})

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

// pnpm add @openai/agents zod
import { Agent, Runner, tool } from "@openai/agents";
import { z } from "zod";

// The working-memory shape lives in TypeScript, not in the prompt.
interface WorkingMemory {
  plan: string[];
  completed: string[];
  notes: string[];
}

const initialMemory: WorkingMemory = {
  plan: [
    "Identify every reader of the global.",
    "Introduce a Config dataclass.",
    "Thread Config through call sites.",
    "Delete the global.",
  ],
  completed: [],
  notes: ["The global is named CONFIG and is set in utils.ts:14."],
};

// Tools receive the typed context and can mutate working memory directly.
const markCompleted = tool({
  name: "mark_completed",
  description: "Mark a plan step as done; moves it from plan to completed.",
  parameters: z.object({ step: z.string() }),
  async execute({ step }, runContext) {
    const mem = runContext.context as WorkingMemory;
    mem.plan = mem.plan.filter((s) => s !== step);
    mem.completed.push(step);
    return `Marked completed: ${step}`;
  },
});

const addNote = tool({
  name: "add_note",
  description: "Pin a fact in working memory.",
  parameters: z.object({ note: z.string() }),
  async execute({ note }, runContext) {
    (runContext.context as WorkingMemory).notes.push(note);
    return `Pinned: ${note}`;
  },
});

// The agent's instructions are dynamic: the working-memory block is rendered
// from the typed context every turn. This is the key move — the prompt always
// shows the latest working memory, no matter how trimmed the conversation
// buffer is.
const refactorAgent = new Agent<WorkingMemory>({
  name: "refactor-agent",
  model: "gpt-5",
  tools: [markCompleted, addNote],
  instructions: (runContext, agent) => {
    const mem = runContext.context;
    return [
      "You refactor TypeScript codebases. Use the working-memory tools to track progress.",
      "",
      "## Working memory (your private scratchpad)",
      "### Plan",
      ...mem.plan.map((s) => `- ${s}`),
      "### Completed",
      ...(mem.completed.length ? mem.completed.map((s) => `- ${s}`) : ["- (none yet)"]),
      "### Pinned notes",
      ...mem.notes.map((n) => `- ${n}`),
    ].join("\n");
  },
});

const result = await Runner.run(refactorAgent, "Continue the refactor.", {
  context: initialMemory,
});

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

# pip install "anthropic>=0.40"
# Minimal notebook substrate using Anthropic's memory tool.
from pathlib import Path
import anthropic

MEMORY_ROOT = Path("./agent_memory").resolve()
MEMORY_ROOT.mkdir(exist_ok=True)

def handle_memory_tool(tool_input: dict) -> str:
    """Client-side handler for the memory_20250818 tool. Path-traversal-safe."""
    cmd = tool_input["command"]
    # Validate every path stays under MEMORY_ROOT — the Anthropic docs flag
    # this as the single most important safeguard.
    def safe_path(raw: str) -> Path:
        p = (MEMORY_ROOT / raw.lstrip("/")).resolve()
        if not p.is_relative_to(MEMORY_ROOT):
            raise ValueError(f"Path escapes memory root: {raw}")
        return p

    if cmd == "view":
        p = safe_path(tool_input["path"])
        if p.is_dir():
            entries = "\n".join(f"{f.stat().st_size}\t/{f.relative_to(MEMORY_ROOT)}"
                                for f in sorted(p.iterdir()))
            return f"Directory listing for {tool_input['path']}:\n{entries}"
        return p.read_text()
    if cmd == "create":
        safe_path(tool_input["path"]).write_text(tool_input["file_text"])
        return f"File created at {tool_input['path']}"
    if cmd == "str_replace":
        p = safe_path(tool_input["path"])
        text = p.read_text()
        if tool_input["old_str"] not in text:
            return f"No replacement: old_str not found in {tool_input['path']}"
        p.write_text(text.replace(tool_input["old_str"], tool_input["new_str"], 1))
        return "The memory file has been edited."
    # Implement insert/delete/rename similarly in production.
    raise ValueError(f"Unknown memory command: {cmd}")

client = anthropic.Anthropic()

# The tool registration is one line; Anthropic injects the working-memory
# protocol into the system prompt automatically.
def run_agent_turn(messages: list[dict]) -> dict:
    return client.messages.create(
        model="claude-opus-4-7",
        max_tokens=2048,
        messages=messages,
        tools=[{"type": "memory_20250818", "name": "memory"}],
    )

# The harness loop dispatches memory tool calls back to the local handler.
# Loop omitted for brevity — see anthropic-sdk-python/examples/memory/basic.py
# for a full runnable version.