Path function gets passed the graphs state_schema instead of the previous nodes returned state

[tgrimminger]

I found out about this while trying to use a private state as described in the documentation here.

My final goal was to have a private state for part of my graph so that I don't have to pollute the main state with additional fields that are only need for a small part of the graph.

Here is a minimum reproducible example, hopefully this makes things more clear:

In this example the special_node returns a StateB (the private state), which the decider is supposed to act on, instead the decider gets a StateA (the graphs state_schema-> the base-state).

from typing_extensions import TypedDict, List
from langgraph.graph import START, END, StateGraph

from langchain_ollama import ChatOllama
from langchain.prompts import PromptTemplate
from langchain.output_parsers import PydanticOutputParser
from langchain.pydantic_v1 import BaseModel, Field
from typing import List

# Define the output structure
class StateA(BaseModel):
    time: str

class StateB(BaseModel):
    test: str

def normal_node(state: StateA):
    state.time = "blub"
    return state

def special_node(state: StateA)->StateB:
    return StateB(test="test")

def decider(state: StateB):
    print(type(state))
    return False

def converter(state: StateB):
    return StateA(time=state.test)

# Graph
workflow = StateGraph(StateA)

# Define the nodes
workflow.add_node(normal_node)
workflow.add_node(special_node)
workflow.add_node(converter)
# Build graph
workflow.set_entry_point("normal_node")
workflow.add_edge("normal_node", "special_node")
workflow.add_conditional_edges(
    "special_node",
    decider,
    {
        True: "converter",
        False: "converter",
    },
)
workflow.add_edge("converter", END)

custom_graph = workflow.compile()
import uuid

def run_agent():
    
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    
    state_dict = custom_graph.invoke(
        {"time":"bla"}, config
    )
    
    return state_dict

result = run_agent()
print(result)
#response

Here I would expect that the decider function gets a StateB, as this is the state that was returned by the last node before it. This is the way it works between two nodes (see documentation). But instead, it gets passed a StateA.

Is this by design (conditional edges can only use the base-state?) or is this a Bug?