[Bug]: prepare_stream regenerate heuristic incorrectly triggers on interrupt resume, destroying interrupt state

[uesleilima]

Pre-flight Checklist

• I have searched existing issues and this hasn't been reported yet.
• I am using the latest version AG-UI.

Describe the Bug

The prepare_stream method in agent.py uses a message-count heuristic to detect "regenerate last response" requests: if the checkpoint has more messages than the frontend sent, it assumes the user wants to time-travel back.

This heuristic is wrong for interrupt resumes. When a LangGraph interrupt() fires, the checkpoint contains an AI message (the tool call that triggered the interrupt) that the frontend never received or displayed. When the frontend later resumes with Command(resume=...), the checkpoint's extra AI message makes prepare_stream think this is a regenerate request.

The result: prepare_stream enters the prepare_regenerate_stream path, which time-travels back and destroys the interrupt state, making it impossible to resume from the interrupt.

Root Cause

agent.py, prepare_stream, line ~460:

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):
    # ... triggers regenerate path

When an interrupt is active:

• Checkpoint has: [HumanMessage, AIMessage(tool_calls=[...])] → 2 messages
• Frontend sends: [HumanMessage] → 1 message (AI tool-call was never shown)
• 2 > 1 → True → regenerate path triggered incorrectly

The has_active_interrupts check at line ~495 comes after the regenerate check, so it never gets a chance to run.

Steps to Reproduce

1. Create a LangGraph agent that uses interrupt() for human-in-the-loop:

from langgraph.types import interrupt

async def approval_node(state):
    answer = interrupt({"question": "Approve?", "options": ["yes", "no"]})
    return {"approved": answer}

2. Send a message that triggers the interrupt.
3. The agent pauses at the interrupt, frontend receives the interrupt event.
4. Frontend sends a resume request with forwarded_props.command.resume = "yes".
5. prepare_stream compares message counts → checkpoint has more → enters regenerate path.
6. Interrupt state is destroyed. The resume value is lost.

Expected Behavior

prepare_stream should detect active interrupts with a resume value before the regenerate heuristic, and bypass the message-count check entirely in that case.

Suggested fix — early return for interrupt resumes:

async def prepare_stream(self, input, agent_state, config):
    ...
    interrupts = self._collect_interrupts(agent_state.tasks)
    has_active_interrupts = len(interrupts) > 0
    resume_input = forwarded_props.get('command', {}).get('resume', None)

    # Interrupt resume must bypass the regenerate heuristic
    if resume_input is not None and has_active_interrupts:
        # Set up state, create Command(resume=...), return stream
        ...

    # Only now check for regenerate
    if len(agent_state.values.get("messages", [])) > len(non_system_messages):
        ...

Environment

ag-ui-langgraph: 0.0.35
ag-ui-protocol:  0.1.18
langgraph:       0.4.7
langchain-core:  0.3.59
Python:          3.13

Related Issues

• LangGraphAgent emits duplicate RUN_STARTED in has_active_interrupts path → RUN_ERROR / INCOMPLETE_STREAM #1584 — duplicate RUN_STARTED in has_active_interrupts path (different root cause, same prepare_stream area)
• 'NoneType' object is not a mapping error in ag-ui-langgraph #702 — NoneType is not a mapping crash in set_message_in_progress (another agent.py bug)
• [Bug]: ag-ui-langgraph drops TOOL_CALL_START/ARGS/END when chat model emits text→tool_call transition in a single chunk #1678 — text→tool_call transition drops TOOL_CALL events (different area)

Additional Context

Workaround: override prepare_stream and detect the interrupt-resume case before the base class's message-count check, then bypass directly to the Command(resume=...) streaming path.

Happy to send a PR with the fix and a regression test if assigned.

[uesleilima]

Adding a related failure mode that the same root-cause fix would solve: the regenerate heuristic also misfires in the opposite direction on dirty checkpoints.

Scenario

1. A stream is cancelled mid-execution (client disconnect, pod eviction, network drop, etc.).
2. The checkpoint is left with an errored task. Often the partial run has committed an AIMessage with tool_calls that were never answered by a ToolMessage.
3. The user edits their last message (same id, new content) and resubmits.
4. prepare_stream runs against the dirty checkpoint:
   • Count check: checkpoint has more messages than the frontend sent (because of the orphan AIMessage) → would trigger regenerate
   • is_continuation check (line ~475): incoming HumanMessage id matches checkpoint id → marks as continuation → SKIPS regenerate
5. Result: stale message content is used, AND on Anthropic, the orphan tool_use causes a 400 because the next request has a broken tool_use/tool_result pairing.

This is the same class of bug as the interrupt-resume false-positive in the original report — the count + id heuristics in prepare_stream are fundamentally fragile when the checkpoint is anything other than the result of a cleanly-completed turn.

Suggested unified fix

Make prepare_stream self-heal dirty checkpoints before applying any heuristics:

async def prepare_stream(self, input, agent_state, config):
    # 1. Self-heal: if checkpoint has errored tasks, fork to last clean parent
    if any(t.error for t in (agent_state.tasks or [])) and agent_state.parent_config:
        parent = await self.graph.aget_state(agent_state.parent_config)
        # Defensive: strip orphan AIMessage(tool_calls=...) without matching ToolMessage
        cleaned = _strip_open_tool_calls(parent.values.get("messages", []))
        await self.graph.aupdate_state(
            agent_state.parent_config,
            {**parent.values, "messages": cleaned},
            as_node=(parent.next[0] if parent.next else "__start__"),
        )
        agent_state = await self.graph.aget_state(config)

    # 2. Interrupt resume (existing fix from this issue)
    if resume_input and has_active_interrupts:
        ...

    # 3. Existing count/id heuristics now operate on a clean state
    ...

This closes both this issue's failure mode (false positive on interrupt resume) AND the dirty-checkpoint failure mode (false negative on edit-after-cancel), because once the checkpoint is clean, add_messages naturally replaces edited content by id and the heuristics make correct decisions.

In our fork we've shipped this as a unified _heal_dirty_checkpoint pre-flight in prepare_stream. Happy to send a PR with the fix + regression tests covering: (a) the interrupt-resume case, (b) the cancelled-stream-then-edit case, (c) the orphan-tool-call case, (d) graceful fallback when there's no parent_config or the checkpointer write fails.