State vs. Memory: Why Context Windows Aren't Enough

In the previous post we built loops: a draft-critique cycle and a tool-calling agent, compiled with InMemorySaver so every super-step was checkpointed under a thread_id. Reuse the thread ID and the agent “remembers” the conversation.

That memory has two hard edges. Start a new thread tomorrow and the agent knows nothing about you: the preference you stated, the bug you fixed together, the architecture decision you argued through. And inside a thread, the only way to give the model information is to put it in the prompt, which fails the moment the information is a 400,000-line monorepo. A bigger context window doesn’t remove either edge; it only moves them and raises the bill.

This post separates the two problems and builds an agent that handles both: short-term state in the checkpointer, long-term user memory in a LangGraph Store, and a codebase in Chroma that the graph searches only when a question needs it.

1. Two problems that look like one

The rows that matter are the read and write paths. Short-term memory is cheap to write and expensive to read: it’s saved automatically, but every token of it is paid for on every model call. Long-term memory is the reverse: writing needs a decision (what is a durable fact?), and reading is cheap because you fetch only what’s relevant.

2. Short-term memory: the thread is the unit

What the checkpointer actually holds

InMemorySaver stores a checkpoint of the full graph state after every super-step, keyed by thread_id. With MessagesState and the add_messages reducer, the message list only grows: every human turn, AI turn, and ToolMessage stays in state for the life of the thread. As the name says, it’s process memory: restart the server and every thread is gone. Production uses PostgresSaver or SqliteSaver; the semantics are the same.

Why “just use a bigger window” doesn’t scale

Current models advertise context windows from 128K to over a million tokens. Three things stand between that number and a usable memory:

• Your serving stack may not give you the window. Hosted APIs reject a request that’s too long, which is at least loud. Self-hosted runtimes often default to a smaller window than the model supports (Ollama, for example, picks one based on available VRAM), and a prompt that doesn’t fit can be cut down without an error ever reaching your code. Know which behavior you have.
• The KV cache is linear in tokens, and large. Take an open-weight model whose architecture is public: Llama 3.1 8B has 32 layers and 8 KV heads of dimension 128 (grouped-query attention). That’s 2 × 32 × 8 × 128 = 65,536 values per token, or 128 KiB per token in fp16. A full 128K window needs about 16 GiB of KV cache on top of the weights, per concurrent sequence.
• Models don’t use long contexts uniformly. Liu et al. (Lost in the Middle) found that accuracy on multi-document QA was highest when the relevant passage sat at the start or end of the context and dropped when it was in the middle. Stuffing more in can make answers worse.

Prefill also costs compute on every call: with full attention it grows quadratically in sequence length, and an agent loop re-sends the whole transcript on every step. Short-term memory needs a budget.

Trim on read, delete on write

There are two places to enforce that budget, and they behave differently:

• On the read path, inside the node, with trim_messages. You send the model the last N tokens, but the checkpoint keeps everything. That’s safe and reversible, and it’s what the code below does.
• On the write path, by returning RemoveMessage(id=...) updates (which add_messages treats as deletions) or by replacing old turns with a running summary. State stays small, but what you drop is gone from this thread unless you saved it somewhere else.

Summaries and deletions are lossy. Before deleting from state, decide whether any of it belongs in long-term memory.

3. Long-term memory: retrieval, not recall

“Long-term memory” covers two different things, and conflating them causes most design mistakes:

1. Corpus memory: a codebase, docs, tickets. Large, mostly read-only, indexed offline by a pipeline, shared across users. This is RAG (Lewis et al., 2020).
2. Agent memory: facts the agent learns while talking to someone (“prefers pytest fixtures”, “service X is being deprecated”). Small, written by the agent at runtime, and scoped to a user or team.

The CoALA framework (Sumers et al., 2023), which LangGraph’s memory docs follow, splits agent memory further into semantic (facts), episodic (past experiences, often stored as few-shot examples), and procedural (rules and instructions, often the system prompt itself). The example below stores semantic memory; the other two use the same machinery with different write policies.

Vector retrieval: what goes wrong in practice

Dense retrieval is approximate nearest-neighbor search over embeddings, usually an HNSW index (Malkov & Yashunin) that trades a little recall for logarithmic-ish query time. The embedding is rarely the hard part. These are:

• Chunking. Split code on syntax boundaries (functions, classes), not fixed character windows, and keep the file path in metadata so the model can cite it. A chunk that cuts a function in half is noise to retrieve.
• Identifiers. Embeddings are weak at exact tokens like get_with_retry or ERR_4012. For code, combine dense retrieval with lexical search (BM25) and merge the rankings; most vector DBs support this hybrid search directly.
• Filters. Constrain by repo, path, language, or tenant before ranking by similarity. Similarity search without metadata filters leaks across tenants and drowns in near-duplicates.
• Evaluation. Treat the retriever as its own model. Build a small set of question → expected-file pairs and track recall@k; most “the LLM hallucinated” bugs in RAG are retrieval misses.

Knowledge graphs: when similarity isn’t the relation you need

“Which services call charge_card, and which of those lack retries?” is a multi-hop question about structure. Nearest-neighbor search returns chunks that look like charge_card; it can’t follow call edges. A graph can: entities as nodes (functions, services, owners), relations as edges, queried by traversal.

For code, build the graph deterministically from static analysis (imports, call graph, ownership files) rather than asking an LLM to extract it: it’s cheaper, complete, and correct. LLM-extracted graphs earn their cost on unstructured text. Microsoft’s GraphRAG (Edge et al., 2024) builds an entity graph over a corpus and pre-summarizes its communities, which helps with global questions (“what are the main themes across these documents?”) that top-k chunk retrieval can’t answer. The trade-off is an expensive indexing step that must be redone when the corpus changes.

A pragmatic default: vectors (plus BM25) for “find me the relevant text”, a graph for “follow the relationships”. Use the graph’s results to pick which chunks to put in the prompt.

4. Wiring both into the graph

The design rule: state holds data, not clients. The checkpointer serializes state after every super-step, so a Chroma client in AgentState would break serialization or end up in every checkpoint. LangGraph gives dependencies two other homes:

• context_schema defines per-run context (user ID, DB clients), passed as invoke(..., context=...) and read in nodes through runtime.context. It isn’t checkpointed.
• compile(store=...) attaches a Store, a namespaced key-value store with optional semantic search, available to nodes as runtime.store and shared across every thread.

What does go in state is the per-turn result of retrieval: the query, the retrieved chunks, the recalled facts. They’re cheap to checkpoint and let you inspect exactly what the model saw on any turn.

flowchart TD
U([User message]) --> recall
recall --> route{route}
route -->|needs code| retrieve
route -->|no code needed| generate
retrieve --> generate
generate --> remember
remember --> A([Answer])
STORE[(Store: memories per user)] -.->|semantic search| recall
STORE <-.-|put new facts| remember
VDB[(Chroma: code chunks)] -.->|top-k chunks| retrieve

Index the codebase (offline)

from pathlib import Path

from langchain_chroma import Chroma
from langchain_text_splitters import Language, RecursiveCharacterTextSplitter

from memory_agent import embeddings

splitter = RecursiveCharacterTextSplitter.from_language(Language.PYTHON, chunk_size=1500, chunk_overlap=150)
codebase = Chroma(collection_name="repo", embedding_function=embeddings, persist_directory=".chroma")

for path in Path("src").rglob("*.py"):
    chunks = splitter.create_documents([path.read_text(encoding="utf-8")], metadatas=[{"source": path.as_posix()}])
    codebase.add_documents(chunks)

from_language(Language.PYTHON) splits on class and def boundaries before falling back to blank lines and characters. In a real pipeline this runs in CI on each merge, keyed by file hash so unchanged files aren’t re-embedded.

The agent

import uuid
from dataclasses import dataclass
from typing import Annotated, TypedDict

from langchain.chat_models import init_chat_model
from langchain.embeddings import init_embeddings
from langchain_chroma import Chroma
from langchain_core.messages import AnyMessage, SystemMessage, trim_messages
from langchain_core.messages.utils import count_tokens_approximately
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.runtime import Runtime
from langgraph.store.memory import InMemoryStore
from pydantic import BaseModel, Field

# Provider-agnostic: e.g. "openai:...", "anthropic:...", "google_genai:...", "ollama:...".
llm = init_chat_model("<provider>:<chat-model>", temperature=0)  # must support structured output
embeddings = init_embeddings("<provider>:<embedding-model>")
EMBED_DIMS = 1536  # must match the embedding model's output size


class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]  # short-term: checkpointed per thread
    query: str             # search query the router wrote ("" = no retrieval)
    documents: list[str]   # code chunks retrieved for *this* turn
    memories: list[str]    # long-term facts recalled for *this* turn


@dataclass
class Context:
    """Per-run dependencies. Not part of state, so never checkpointed."""
    user_id: str
    codebase: Chroma


class Route(BaseModel):
    """Decide whether the question needs the code index."""
    needs_code: bool = Field(description="True if answering requires reading this repository's code")
    search_query: str = Field(description="Query for the code index, or empty if needs_code is false")


class Facts(BaseModel):
    """Durable facts about the user worth remembering across sessions."""
    facts: list[str] = Field(description="Preferences, decisions, or project facts. Empty if none.")


def memory_ns(user_id: str) -> tuple[str, ...]:
    return ("users", user_id, "memories")


def recall(state: AgentState, runtime: Runtime[Context]):
    question = state["messages"][-1].content
    hits = runtime.store.search(memory_ns(runtime.context.user_id), query=question, limit=3)
    return {"memories": [h.value["text"] for h in hits]}


def route(state: AgentState):
    decision = llm.with_structured_output(Route).invoke(state["messages"][-1:])
    # Clear last turn's retrieval: state persists across turns in a thread.
    return {"query": decision.search_query if decision.needs_code else "", "documents": []}


def retrieve(state: AgentState, runtime: Runtime[Context]):
    docs = runtime.context.codebase.similarity_search(state["query"], k=4)
    return {"documents": [f"# {d.metadata['source']}\n{d.page_content}" for d in docs]}


def generate(state: AgentState):
    system = "You are a senior engineer helping with this repository. Be concise."
    if state["memories"]:
        system += "\n\nWhat you know about this user:\n" + "\n".join(f"- {m}" for m in state["memories"])
    if state["documents"]:
        system += "\n\nRelevant code:\n\n" + "\n\n".join(state["documents"])
    # Bound the history we *send*; the checkpoint still keeps every message.
    history = trim_messages(
        state["messages"],
        max_tokens=4000,
        token_counter=count_tokens_approximately,
        strategy="last",
        start_on="human",
    )
    return {"messages": [llm.invoke([SystemMessage(system), *history])]}


def remember(state: AgentState, runtime: Runtime[Context]):
    extracted = llm.with_structured_output(Facts).invoke(
        [SystemMessage("Extract durable facts about the user from their last message.")]
        + state["messages"][-2:-1]
    )
    for fact in extracted.facts:
        runtime.store.put(memory_ns(runtime.context.user_id), str(uuid.uuid4()), {"text": fact})
    return {}


def needs_retrieval(state: AgentState) -> str:
    return "retrieve" if state["query"] else "generate"


builder = StateGraph(AgentState, context_schema=Context)
builder.add_node("recall", recall)
builder.add_node("route", route)
builder.add_node("retrieve", retrieve)
builder.add_node("generate", generate)
builder.add_node("remember", remember)
builder.add_edge(START, "recall")
builder.add_edge("recall", "route")
builder.add_conditional_edges("route", needs_retrieval, ["retrieve", "generate"])
builder.add_edge("retrieve", "generate")
builder.add_edge("generate", "remember")
builder.add_edge("remember", END)

app = builder.compile(
    checkpointer=InMemorySaver(),  # short-term: one thread's state
    store=InMemoryStore(index={"embed": embeddings, "dims": EMBED_DIMS, "fields": ["text"]}),  # long-term
)

Node by node:

• recall runs a semantic search over the user’s namespace in the Store. The namespace comes from runtime.context, which your server fills from the authenticated session, never from model output.
• route asks the model for a structured Route. This is the “retrieve before reasoning” decision, made explicitly in its own node so you can log it, evaluate it, and override it. It also resets documents, because a plain channel keeps its last value across turns: without the reset, a small-talk follow-up would still carry the previous turn’s code.
• retrieve queries Chroma through the injected client and writes plain strings into state.
• generate assembles the prompt from three sources with three budgets: recalled facts (at most 3), code chunks (k = 4), and the trimmed thread history (about 4,000 tokens).
• remember extracts durable facts from the user’s last message and writes them to the Store, where any future thread for this user can find them.

Running it across two sessions

from langchain_chroma import Chroma

from memory_agent import Context, app, embeddings

codebase = Chroma(collection_name="repo", embedding_function=embeddings, persist_directory=".chroma")
ctx = Context(user_id="prabesh", codebase=codebase)

monday = {"configurable": {"thread_id": "monday"}}
app.invoke({"messages": [("user", "I always use pytest fixtures, never unittest classes.")]}, monday, context=ctx)

tuesday = {"configurable": {"thread_id": "tuesday"}}  # new thread: empty short-term state
out = app.invoke({"messages": [("user", "Write a test for the retry logic in our HTTP client.")]}, tuesday, context=ctx)

print(out["memories"])                                   # recalled from Monday, via the Store
print([d.splitlines()[0] for d in out["documents"]])     # retrieved from Chroma
print(len(app.get_state(monday).values["messages"]))     # Monday's thread is untouched

With a scripted chat model (so the routing and extraction are deterministic), a real embedding model, and a three-file Chroma index, the run printed:

['Prefers pytest fixtures over unittest classes']
['# src/http/client.py', '# src/db/users.py', '# src/config/logging.py']
2

The Tuesday thread started empty, yet the prompt for generate contained both the Monday preference and src/http/client.py. A follow-up “Thanks, that’s all.” on the same thread routed straight to generate with documents == [].

sequenceDiagram
participant U as User
participant G as Graph
participant C as Checkpointer
participant S as Store
Note over U,S: Monday, thread_id=monday
U->>G: I always use pytest fixtures
G->>S: put users/prabesh/memories
G->>C: save monday checkpoints
Note over U,S: Tuesday, thread_id=tuesday
U->>G: Write a test for the retry logic
G->>C: load tuesday, nothing there yet
G->>S: search users/prabesh/memories
S-->>G: Prefers pytest fixtures
G-->>U: Test written with fixtures

5. Production notes

• Swap the in-memory backends. PostgresSaver for checkpoints and PostgresStore (pgvector) for memories, both from langgraph-checkpoint-postgres. Chroma can run as a server, or swap in Pinecone, pgvector, or whatever your platform already operates; only the Context field type changes.
• Move memory writes off the hot path. remember adds a full LLM call to every turn. LangGraph’s memory docs describe writing memories in the background instead: queue the transcript and extract after the response is sent, or batch it when the thread goes idle.
• Consolidate, don’t just append. Writing a new UUID per fact means “prefers pytest” gets stored twenty times, and contradicted facts never get deleted. Search before writing and update the existing item, or use a library like LangMem that manages consolidation for you.
• Treat memory as an attack surface. Anything a user (or a retrieved document) says can become a “fact” that is injected into every future system prompt. Validate what gets stored, keep memories per tenant, and give users a way to view and delete them.
• Expire things. Store.put takes a ttl in minutes on stores that support it: PostgresStore does, InMemoryStore raises NotImplementedError. Project facts go stale; a memory with no expiry policy eventually becomes wrong.
• Re-index on change. A code index that lags main by a week answers confidently about code that no longer exists. Store the commit SHA in chunk metadata and show it in citations.

6. Pitfalls

• Putting clients or file handles in state. They end up in every checkpoint, or break serialization outright. Use context_schema and runtime.context.
• Forgetting that state persists across turns. Overwrite channels like documents keep their last value. Reset per-turn fields explicitly.
• Trimming tool-call pairs apart. If trimming keeps an AIMessage with tool_calls but drops its ToolMessage (or the reverse), most providers reject the request. Trim on human-turn boundaries (start_on="human") and test with tool-heavy threads.
• Using the thread as long-term memory. Reusing one thread_id forever turns state into an unbounded log that you then have to trim anyway. Threads are for conversations; the Store is for what outlives them.
• Trusting the default context length. Check what window your serving stack actually gives you, and log prompt token counts so you notice when a turn approaches it.
• Evaluating only the final answer. Log query, documents, and memories per turn (they’re in state, so get_state_history already has them) and score retrieval separately.

Study checklist

1. Run memory_agent.py against a real chat model and a real repo; try two providers and compare the routing. Print app.get_state(config).values["documents"] after each turn and judge the retrievals yourself before reading the answers.
2. Remove the "documents": [] reset in route, ask a code question and then “thanks”, and watch stale code reach the prompt.
3. Replace trim_messages with a summarization node that writes a running summary and returns RemoveMessages for the turns it summarized. Compare token counts per call over a 30-turn thread.
4. Make remember search the namespace first and update a near-duplicate instead of inserting.
5. Add BM25 over the same chunks, merge the rankings with reciprocal rank fusion, and measure recall@4 on 20 hand-labeled questions against dense-only retrieval.
6. Swap InMemorySaver and InMemoryStore for the Postgres versions and restart the process between Monday and Tuesday.

References