fix(langgraph): inject runtime with store for conditional edges in update_state operations #6340

[DraPraks]

Description

Fixes #6340

This PR resolves a bug where conditional edge functions with a store: BaseStore parameter would fail with ValueError: Missing required config key 'store' when invoked during update_state() or bulk_update_state() operations.

Problem

When using conditional edges with dependency injection for the store parameter, the graph would work correctly during normal execution (invoke(), stream()), but fail when state was updated via update_state() or bulk_update_state(). This is because the runtime configuration containing the store was not being passed to the writer functions (conditional edges) in the bulk update code paths.

Example that was failing:

def route_with_store(
    state: State, 
    config: RunnableConfig, 
    store: BaseStore  # <-- This caused the error
) -> Literal["next_node", END]:
    # Access store for routing logic
    preferences = list(store.search(namespace))
    return "next_node" if preferences else END

builder.add_conditional_edges("node", route_with_store)
graph = builder.compile(checkpointer=checkpointer, store=store)

# This worked fine
graph.invoke({"value": 1}, config)

# This failed with ValueError: Missing required config key 'store'
graph.update_state(config, {"value": 2}, as_node="node")

Solution

Applied the same runtime setup pattern used in invoke() and stream() methods to both bulk_update_state() and abulk_update_state() methods:

1. Extract store from config or use self.store
2. Create Runtime object with store, context, stream_writer, and previous state
3. Merge with parent runtime
4. Pass runtime via CONFIG_KEY_RUNTIME in configurable dict when invoking writers

This ensures conditional edges have access to all dependency-injected parameters during state updates, maintaining consistency across all execution paths.

Changes

Core Fix (libs/langgraph/langgraph/pregel/main.py)

• Added PREVIOUS constant import
• Added runtime setup in bulk_update_state() (lines ~1794-1808)
• Added runtime setup in abulk_update_state() (lines ~2247-2261)
• Pass runtime via CONFIG_KEY_RUNTIME to writer invocations in both methods

Tests (libs/langgraph/tests/test_pregel.py & test_pregel_async.py)

Added 8 test cases (4 sync + 4 async):

1. test_conditional_edge_with_store_in_update_state - Basic regression test for the reported bug
2. test_conditional_edge_with_store_bulk_update_state - Tests bulk update scenario with multiple state updates
3. test_conditional_edge_store_with_config_runtime - Tests runtime merging when parent config already has runtime
4. test_conditional_edge_without_store_still_works - Backward compatibility test

Verifying:

• Conditional edges can access store parameter without errors
• Store operations (get, put, search) work correctly
• Routing logic executes properly
• Graph execution continues correctly after state updates

Checklist

• [x] Bug fix (non-breaking change which fixes an issue)
• [x] Includes tests that fail before the fix and pass after
• [x] Follows existing code patterns and style
• [x] Passed make format, make lint, make test

---

Notes

• Scope: libs/langgraph
• Branch: fix/store-conditional-edge-update-state-6340
• The fix mirrors the existing pattern used in invoke() and stream() methods for consistency
• All changes are isolated to the Pregel class in main.py

[vercel[bot]]

@DraPraks is attempting to deploy a commit to the LangChain Team on Vercel.

A member of the Team first needs to authorize it.