ValueError: Message ID not found in history when resuming threads in ag-ui-langgraph 0.0.21

Description

Package: ag-ui-langgraph
Version: 0.0.21
Severity: Critical - Blocks thread resumption functionality

Summary

When resuming an existing thread (e.g., after page refresh), the prepare_stream()
method in LangGraphAgent crashes with ValueError: Message ID not found in history.
This occurs because the method incorrectly treats thread resumption as a
"regenerate" request and tries to find client-generated message IDs in checkpoint
history.

---
Steps to Reproduce

1. Setup: LangGraph agent with AsyncPostgresSaver checkpointer
2. First interaction: Send a message, receive response (checkpoint saved)
3. Reload page: Frontend refreshes, maintaining thread_id via localStorage
4. Resume thread: Send a new message with the same thread_id
5. Crash: Backend crashes with ValueError: Message ID not found in history

Minimal Code Example

from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langchain_anthropic import ChatAnthropic
from langgraph.graph import StateGraph, START, END
from ag_ui_langgraph import LangGraphAgent

# Create graph with checkpointer
checkpointer = AsyncPostgresSaver(...)
graph = workflow.compile(checkpointer=checkpointer)

# Create AG-UI agent
agent = LangGraphAgent(
name="test_agent",
graph=graph,
)

# First invocation (works)
config = {"configurable": {"thread_id": "thread-123"}}
result1 = await agent.run(input1) # ✅ Works

# Second invocation (crashes)
result2 = await agent.run(input2) # ❌ Crashes with ValueError

---
Error Stack Trace

ValueError: Message ID not found in history
File "ag_ui_langgraph/agent.py", line 297, in prepare_stream
return await self.prepare_regenerate_stream(...)
File "ag_ui_langgraph/agent.py", line 369, in prepare_regenerate_stream
time_travel_checkpoint = await
self.get_checkpoint_before_message(message_checkpoint.id, thread_id)
File "ag_ui_langgraph/agent.py", line 843, in get_checkpoint_before_message
raise ValueError("Message ID not found in history")

---
Root Cause Analysis

The Bug (lines 287-301 in agent.py)

non_system_messages = [msg for msg in langchain_messages if not isinstance(msg,
SystemMessage)]
if len(agent_state.values.get("messages", [])) > len(non_system_messages):
# Find the last user message by working backwards from the last message
last_user_message = None
for i in range(len(langchain_messages) - 1, -1, -1):
if isinstance(langchain_messages[i], HumanMessage):
last_user_message = langchain_messages[i]
break

if last_user_message:
return await self.prepare_regenerate_stream( # ❌ CRASHES HERE
input=input,
message_checkpoint=last_user_message,
config=config
)

Why It Fails

1. Scenario: Thread resumption after page reload
- Checkpoint contains: 4 messages (with IDs: msg-1, msg-2, msg-3, msg-4)
- New request contains: 1 message (with NEW client-generated ID: msg-5)
2. Faulty Logic: The condition len(checkpoint_messages) > len(incoming_messages)
triggers
- Code assumes: "More messages in checkpoint → must be regenerate request"
- Reality: This is a normal thread resumption with a NEW message
3. The Crash: prepare_regenerate_stream() tries to find msg-5 in checkpoint
history
- Expected: Message ID exists in checkpoint (for regenerate)
- Actual: Message ID is NEW (client-generated for thread resumption)
- Result: ValueError: Message ID not found in history

---
Expected Behavior

Thread resumption should work seamlessly:
- Load checkpoint for existing thread_id
- Accept new message with new ID
- Continue conversation from checkpoint state

---
Actual Behavior

- Crash with ValueError: Message ID not found in history
- Thread resumption impossible
- Forces users to create new threads for every page refresh

---
Proposed Fix

Add message ID validation before calling prepare_regenerate_stream():

non_system_messages = [msg for msg in langchain_messages if not isinstance(msg,
SystemMessage)]
if len(agent_state.values.get("messages", [])) > len(non_system_messages):
last_user_message = None
for i in range(len(langchain_messages) - 1, -1, -1):
if isinstance(langchain_messages[i], HumanMessage):
last_user_message = langchain_messages[i]
break

if last_user_message:
# ✅ FIX: Check if message ID exists in checkpoint before regenerating
checkpoint_message_ids = {
getattr(msg, 'id', None)
for msg in agent_state.values.get("messages", [])
}

# Only regenerate if message ID is found in history
if last_user_message.id in checkpoint_message_ids:
return await self.prepare_regenerate_stream(
input=input,
message_checkpoint=last_user_message,
config=config
)
# Otherwise, continue with normal thread resumption (fall through)

---
Workaround

Override prepare_stream() in a custom agent class:

class FixedLangGraphAgent(LangGraphAgent):
async def prepare_stream(self, input, agent_state, config):
# ... copy parent logic with fix above

See full workaround implementation: [link to your repo if you want to share]

---
Environment

- ag-ui-langgraph: 0.0.21
- langgraph: 1.0.3
- langchain: 1.0.7
- langgraph-checkpoint-postgres: 2.0.0
- Python: 3.12

---
Additional Context

This bug affects all thread resumption scenarios including:
- Page refreshes in web apps
- Session restoration after browser restart
- Long-running conversations with checkpoints
- Any scenario where thread_id is reused with new messages

The regenerate feature itself is valuable, but it needs proper detection logic to
distinguish between:
- Regenerate request: User wants to regenerate an existing message (message ID in
checkpoint)
- Thread resumption: User sends new message on existing thread (message ID NOT in
checkpoint)

---
Impact

- Critical: Breaks core thread persistence functionality
- Workaround: Possible but requires custom agent class
- Users Affected: Anyone using ag-ui-langgraph with LangGraph checkpointing