Error when using pydantic model in Command.update

[vbarda]

Privileged issue

• I am a LangChain maintainer, or was asked directly by a LangChain maintainer to create an issue here.

Issue Content

from pydantic import BaseModel
from langgraph.graph import StateGraph, START
from langgraph.types import Command

class State(BaseModel):
    foo: str

def node_a(state: State):
    return Command(goto="node_b", update=State(foo=state.foo + "bar"))

def node_b(state: State):
    return State(foo=state.foo + "baz")

builder = StateGraph(State)
builder.add_node(node_a)
builder.add_node(node_b)
builder.add_edge(START, "node_a")
graph = builder.compile()

graph.invoke(State(foo=""))

Error:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[26], line 20
     17 builder.add_edge(START, "node_a")
     18 graph = builder.compile()
---> 20 graph.invoke(State(foo=""))

File [~/development/langgraph/libs/langgraph/langgraph/pregel/__init__.py:1936](http://localhost:8888/lab/tree/~/development/langgraph/libs/langgraph/langgraph/pregel/__init__.py#line=1935), in Pregel.invoke(self, input, config, stream_mode, output_keys, interrupt_before, interrupt_after, debug, **kwargs)
   1934 else:
   1935     chunks = []
-> 1936 for chunk in self.stream(
   1937     input,
   1938     config,
   1939     stream_mode=stream_mode,
   1940     output_keys=output_keys,
   1941     interrupt_before=interrupt_before,
   1942     interrupt_after=interrupt_after,
   1943     debug=debug,
   1944     **kwargs,
   1945 ):
   1946     if stream_mode == "values":
   1947         latest = chunk

File [~/development/langgraph/libs/langgraph/langgraph/pregel/__init__.py:1655](http://localhost:8888/lab/tree/~/development/langgraph/libs/langgraph/langgraph/pregel/__init__.py#line=1654), in Pregel.stream(self, input, config, stream_mode, output_keys, interrupt_before, interrupt_after, debug, subgraphs)
   1649     get_waiter = None  # type: ignore[assignment]
   1650 # Similarly to Bulk Synchronous Parallel [/](http://localhost:8888/) Pregel model
   1651 # computation proceeds in steps, while there are channel updates
   1652 # channel updates from step N are only visible in step N+1
   1653 # channels are guaranteed to be immutable for the duration of the step,
   1654 # with channel updates applied only at the transition between steps
-> 1655 while loop.tick(input_keys=self.input_channels):
   1656     for _ in runner.tick(
   1657         loop.tasks.values(),
   1658         timeout=self.step_timeout,
   (...)
   1661     ):
   1662         # emit output
   1663         yield from output()

File [~/development/langgraph/libs/langgraph/langgraph/pregel/loop.py:394](http://localhost:8888/lab/tree/~/development/langgraph/libs/langgraph/langgraph/pregel/loop.py#line=393), in PregelLoop.tick(self, input_keys)
    392 # apply writes to managed values
    393 for key, values in mv_writes.items():
--> 394     self._update_mv(key, values)
    395 # produce values output
    396 self._emit(
    397     "values", map_output_values, self.output_keys, writes, self.channels
    398 )

File [~/development/langgraph/libs/langgraph/langgraph/pregel/loop.py:830](http://localhost:8888/lab/tree/~/development/langgraph/libs/langgraph/langgraph/pregel/loop.py#line=829), in SyncPregelLoop._update_mv(self, key, values)
    829 def _update_mv(self, key: str, values: Sequence[Any]) -> None:
--> 830     return self.submit(cast(WritableManagedValue, self.managed[key]).update, values)

KeyError: '__root__'