more

langgraph-go/advancedgraph/graph.go

@@ -217,7 +217,10 @@ func compileNodeExecutor(fn any, expectedStateType reflect.Type) (nodeExecutor,
 		if err != nil {
 			return Command{}, fmt.Errorf("node `%s` state decode failed: %w", NodeName(fn), err)
 		}
-		beforeState := stateArg.Interface()
+		beforeStateMap, ok := structValueToMap(stateArg)
+		if !ok {
+			return Command{}, fmt.Errorf("node `%s` failed to snapshot state fields", NodeName(fn))
+		}
 		args := []reflect.Value{
 			reflect.ValueOf(ctx),
 			reflect.Zero(inputType),
@@ -251,7 +254,7 @@ func compileNodeExecutor(fn any, expectedStateType reflect.Type) (nodeExecutor,
 				)
 			}
 		}
-		cmd.Update = reduceStructUpdateToChangedFields(beforeState, cmd.Update, stateType)
+		cmd.Update = reduceStructUpdateToChangedFields(beforeStateMap, cmd.Update, stateType)
 		if out[1].IsNil() {
 			return cmd, nil
 		}
@@ -306,38 +309,45 @@ func mustTypeOf[T any]() reflect.Type {
 	return reflect.TypeOf((*T)(nil)).Elem()
 }
 
-func reduceStructUpdateToChangedFields(before any, update any, stateType reflect.Type) any {
+func reduceStructUpdateToChangedFields(
+	before map[string]any,
+	update any,
+	stateType reflect.Type,
+) any {
 	if update == nil {
 		return nil
 	}
-	beforeV := reflect.ValueOf(before)
 	updateV := reflect.ValueOf(update)
-	if !beforeV.IsValid() || !updateV.IsValid() {
-		return update
+	if !updateV.IsValid() {
+		panic("internal invariant violated: update is non-nil but reflect value is invalid")
 	}
-	if beforeV.Type() != stateType || updateV.Type() != stateType {
-		return update
+	if updateV.Type() != stateType {
+		panic(fmt.Sprintf(
+			"internal invariant violated: update type mismatch in reducer: got %s, expected %s",
+			updateV.Type().String(),
+			stateType.String(),
+		))
 	}
 	if stateType.Kind() != reflect.Struct {
-		return update
+		panic(fmt.Sprintf(
+			"internal invariant violated: stateType must be struct in reducer, got %s",
+			stateType.Kind().String(),
+		))
+	}
+	updateMap, ok := structValueToMap(updateV)
+	if !ok {
+		panic(fmt.Sprintf(
+			"internal invariant violated: failed to convert struct update to map for type %s",
+			stateType.String(),
+		))
 	}
-
 	changed := make(map[string]any)
-	for i := 0; i < stateType.NumField(); i++ {
-		field := stateType.Field(i)
-		if field.PkgPath != "" {
+	for key, updateValue := range updateMap {
+		prevValue, ok := before[key]
+		if ok && reflect.DeepEqual(prevValue, updateValue) {
 			continue
 		}
-		beforeField := beforeV.Field(i)
-		updateField := updateV.Field(i)
-		if reflect.DeepEqual(beforeField.Interface(), updateField.Interface()) {
-			continue
-		}
-		key := jsonFieldName(field)
-		if key == "-" {
-			continue
-		}
-		changed[key] = updateField.Interface()
+		changed[key] = updateValue
 	}
 	if len(changed) == 0 {
 		return nil
@@ -345,14 +355,49 @@ func reduceStructUpdateToChangedFields(before any, update any, stateType reflect
 	return changed
 }
 
-func jsonFieldName(field reflect.StructField) string {
+func structValueToMap(value reflect.Value) (map[string]any, bool) {
+	if !value.IsValid() || value.Kind() != reflect.Struct {
+		return nil, false
+	}
+	out := make(map[string]any, value.NumField())
+	typ := value.Type()
+	for i := 0; i < value.NumField(); i++ {
+		field := typ.Field(i)
+		if field.PkgPath != "" {
+			continue
+		}
+		name, omitEmpty, skip := parseJSONFieldTag(field)
+		if skip {
+			continue
+		}
+		fv := value.Field(i)
+		if omitEmpty && fv.IsZero() {
+			continue
+		}
+		out[name] = fv.Interface()
+	}
+	return out, true
+}
+
+func parseJSONFieldTag(field reflect.StructField) (name string, omitEmpty bool, skip bool) {
 	tag := field.Tag.Get("json")
+	if tag == "-" {
+		return "", false, true
+	}
 	if tag == "" {
-		return field.Name
+		return field.Name, false, false
 	}
 	parts := strings.Split(tag, ",")
-	if len(parts) == 0 || parts[0] == "" {
-		return field.Name
+	fieldName := parts[0]
+	if fieldName == "" {
+		fieldName = field.Name
 	}
-	return parts[0]
+	omit := false
+	for _, opt := range parts[1:] {
+		if opt == "omitempty" {
+			omit = true
+			break
+		}
+	}
+	return fieldName, omit, false
 }

langgraph-go/tests/test_update_elision_test.go

@@ -1,6 +1,7 @@
 package tests
 
 import (
+	"reflect"
 	"testing"
 	"time"
 
@@ -8,39 +9,103 @@ import (
 )
 
 type updateElisionState struct {
-	X int `json:"x"`
+	X  int             `json:"x"`
+	S  updateStruct    `json:"s"`
+	M  map[string]int  `json:"m"`
+	L  []int           `json:"l"`
+	PS *updateStruct   `json:"ps"`
+	PM *map[string]int `json:"pm"`
+	PL *[]int          `json:"pl"`
+}
+
+type updateStruct struct {
+	V int `json:"v"`
 }
 
 type updateElisionWorkflow struct{}
 
-func (w *updateElisionWorkflow) startNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
+func makeState(v int) updateElisionState {
+	m := map[string]int{"n": v}
+	l := []int{v}
+	return updateElisionState{
+		X:  v,
+		S:  updateStruct{V: v},
+		M:  map[string]int{"n": v},
+		L:  []int{v},
+		PS: &updateStruct{V: v},
+		PM: &m,
+		PL: &l,
+	}
+}
+
+func (w *updateElisionWorkflow) startNoopNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
 	return ag.Command{
 		Goto: []ag.Send{
 			{Node: w.fastNode},
-			{Node: w.slowNode},
+			{Node: w.slowNoopNode},
+		},
+	}, nil
+}
+
+func (w *updateElisionWorkflow) startChangedNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
+	return ag.Command{
+		Goto: []ag.Send{
+			{Node: w.fastNode},
+			{Node: w.slowChangedNode},
 		},
 	}, nil
 }
 
 func (w *updateElisionWorkflow) fastNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) {
-	state.X = 1
+	_ = state
+	return ag.Command{Update: makeState(1)}, nil
+}
+
+func (w *updateElisionWorkflow) slowNoopNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) {
+	time.Sleep(100 * time.Millisecond)
+	// Returns same state as initial snapshot; SDK should elide this update.
 	return ag.Command{Update: state}, nil
 }
 
-func (w *updateElisionWorkflow) slowNode(ctx *ag.Context, _ any, state updateElisionState) (ag.Command, error) {
+func (w *updateElisionWorkflow) slowChangedNode(ctx *ag.Context, _ any, _ updateElisionState) (ag.Command, error) {
 	time.Sleep(100 * time.Millisecond)
-	// Returns same X as initial snapshot; SDK should elide this update.
-	return ag.Command{Update: state}, nil
+	// Real change should not be elided.
+	return ag.Command{Update: makeState(2)}, nil
+}
+
+func assertStateEquals(t *testing.T, got updateElisionState, expected updateElisionState) {
+	t.Helper()
+	if got.X != expected.X {
+		t.Fatalf("unexpected X: got=%d want=%d", got.X, expected.X)
+	}
+	if got.S != expected.S {
+		t.Fatalf("unexpected S: got=%#v want=%#v", got.S, expected.S)
+	}
+	if !reflect.DeepEqual(got.M, expected.M) {
+		t.Fatalf("unexpected M: got=%#v want=%#v", got.M, expected.M)
+	}
+	if !reflect.DeepEqual(got.L, expected.L) {
+		t.Fatalf("unexpected L: got=%#v want=%#v", got.L, expected.L)
+	}
+	if got.PS == nil || expected.PS == nil || *got.PS != *expected.PS {
+		t.Fatalf("unexpected PS: got=%#v want=%#v", got.PS, expected.PS)
+	}
+	if got.PM == nil || expected.PM == nil || !reflect.DeepEqual(*got.PM, *expected.PM) {
+		t.Fatalf("unexpected PM: got=%#v want=%#v", got.PM, expected.PM)
+	}
+	if got.PL == nil || expected.PL == nil || !reflect.DeepEqual(*got.PL, *expected.PL) {
+		t.Fatalf("unexpected PL: got=%#v want=%#v", got.PL, expected.PL)
+	}
 }
 
 func TestNoopSlowUpdateDoesNotOverrideFastUpdate(t *testing.T) {
 	workflow := &updateElisionWorkflow{}
 	graph := ag.NewAdvancedStateGraph[updateElisionState]()
-	graph.AddEntryNode(workflow.startNode)
+	graph.AddEntryNode(workflow.startNoopNode)
 	graph.AddNode(workflow.fastNode)
-	graph.AddFinishNode(workflow.slowNode)
+	graph.AddFinishNode(workflow.slowNoopNode)
 
-	handler, err := graph.Compile().Start(nil, updateElisionState{X: 0})
+	handler, err := graph.Compile().Start(nil, makeState(0))
 	if err != nil {
 		t.Fatalf("start failed: %v", err)
 	}
@@ -48,7 +113,23 @@ func TestNoopSlowUpdateDoesNotOverrideFastUpdate(t *testing.T) {
 	if err != nil {
 		t.Fatalf("result failed: %v", err)
 	}
-	if result.X != 1 {
-		t.Fatalf("expected fast update to win, got x=%d", result.X)
-	}
+	assertStateEquals(t, result, makeState(1))
+}
+
+func TestChangedSlowUpdateOverridesFastUpdate(t *testing.T) {
+	workflow := &updateElisionWorkflow{}
+	graph := ag.NewAdvancedStateGraph[updateElisionState]()
+	graph.AddEntryNode(workflow.startChangedNode)
+	graph.AddNode(workflow.fastNode)
+	graph.AddFinishNode(workflow.slowChangedNode)
+
+	handler, err := graph.Compile().Start(nil, makeState(0))
+	if err != nil {
+		t.Fatalf("start failed: %v", err)
+	}
+	result, err := handler.WaitForResult()
+	if err != nil {
+		t.Fatalf("result failed: %v", err)
+	}
+	assertStateEquals(t, result, makeState(2))
 }

libs/langgraph/langgraph/advanced_graph/state.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import copy
 import inspect
 from collections.abc import Callable, Coroutine, Sequence
 from dataclasses import dataclass
@@ -252,7 +253,7 @@ class _GraphEngineRun:
     def _execute_node_for_rust(
         self, node_name: str, node_input: Any, state: Any
     ) -> dict[str, Any]:
-        before_state_markers = _state_shallow_markers(state)
+        before_state_snapshot = copy.deepcopy(state) if isinstance(state, dict) else None
 
         if node_name not in self._nodes:
             raise ValueError(f"Unknown node `{node_name}`")
@@ -270,11 +271,9 @@ class _GraphEngineRun:
 
         if update is None and isinstance(state, dict):
             # Preserve in-place state mutations for prototype nodes like wait_node.
-            if _has_shallow_state_change(before_state_markers, state):
-                update = state
-        elif isinstance(update, dict):
-            if not _has_shallow_update_change(before_state_markers, update):
-                update = None
+            update = state
+        if isinstance(update, dict):
+            update = _reduce_update_to_changed_fields(before_state_snapshot, update)
 
         return {
             "update": update,
@@ -390,44 +389,16 @@ def _resolve_target_name(target: Any) -> str:
     raise ValueError(f"Unsupported node target type: {type(target)!r}")
 
 
-def _state_shallow_markers(state: Any) -> dict[str, Any] | None:
-    if not isinstance(state, dict):
-        return None
-    return {k: _value_shallow_marker(v) for k, v in state.items()}
-
-
-def _value_shallow_marker(value: Any) -> Any:
-    if value is None or isinstance(value, (bool, int, float, str, bytes)):
-        return ("primitive", value)
-    if isinstance(value, (list, tuple, set, dict)):
-        return ("container", type(value).__name__, id(value), len(value))
-    return ("object", type(value).__name__, id(value))
-
-
-def _has_shallow_state_change(before: dict[str, Any] | None, state: Any) -> bool:
-    if before is None or not isinstance(state, dict):
-        return True
-    if len(before) != len(state):
-        return True
-    for key, old_marker in before.items():
-        if key not in state:
-            return True
-        if old_marker != _value_shallow_marker(state[key]):
-            return True
-    return False
-
-
-def _has_shallow_update_change(
+def _reduce_update_to_changed_fields(
     before: dict[str, Any] | None, update: dict[str, Any]
-) -> bool:
+) -> dict[str, Any] | None:
     if before is None:
-        return True
+        return update
+    changed: dict[str, Any] = {}
     for key, new_value in update.items():
-        if key not in before:
-            return True
-        if _value_shallow_marker(new_value) != before[key]:
-            return True
-    return False
+        if key not in before or before[key] != new_value:
+            changed[key] = new_value
+    return changed or None
 
 
 def _invoke_node(node: Callable[..., Any], ctx: Context, node_input: Any, state: Any) -> Any:

libs/langgraph/tests/advanced-graph/test_update_elision.py

@@ -1,35 +1,125 @@
 import asyncio
+from dataclasses import dataclass
+from pydantic import BaseModel
 from typing_extensions import TypedDict
 
 import pytest
 
-from langgraph.advanced_graph import AdvancedStateGraph
+from langgraph.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
 from langgraph.types import Command, Send
 
 pytestmark = pytest.mark.anyio
 
 
+@dataclass(frozen=True)
+class DataClassPayload:
+    value: int
+
+
+class PydanticPayload(BaseModel):
+    value: int
+
+
+class InnerTypedDict(TypedDict):
+    flag: bool
+    n: int
+
+
 class UpdateElisionState(TypedDict):
     x: int
+    dc: DataClassPayload
+    model: PydanticPayload
+    td: InnerTypedDict
+    obj: dict[str, int]
+    items: list[int]
+
+
+def _initial_state() -> UpdateElisionState:
+    return {
+        "x": 0,
+        "dc": DataClassPayload(0),
+        "model": PydanticPayload(value=0),
+        "td": {"flag": False, "n": 0},
+        "obj": {"n": 0},
+        "items": [0],
+    }
+
 
 
 async def test_noop_slow_update_does_not_override_fast_update() -> None:
-    graph = AdvancedStateGraph(UpdateElisionState)
+    graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
 
     async def start_node(state: UpdateElisionState) -> Command:
         return Command(goto=[Send("fast_node", None), Send("slow_node", None)])
 
-    async def fast_node(state: UpdateElisionState) -> dict[str, int]:
-        return {"x": 1}
+    async def fast_node(state: UpdateElisionState) -> UpdateElisionState:
+        state.x = 1
+        state.dc.value = 1
+        state.model.value = 1
+        state.td["flag"] = True
+        state.td["n"] = 1
+        state.obj["n"] = 1
+        state.items.append(1)
+        return state
 
-    async def slow_node(state: UpdateElisionState) -> dict[str, int]:
+    async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
         await asyncio.sleep(0.1)
-        # Returns the same value as initial snapshot.
-        return {"x": state["x"]}
+        # Returns the same values as the initial snapshot.
+        return state
 
     graph.add_entry_node(start_node)
     graph.add_node(fast_node)
     graph.add_finish_node(slow_node)
 
-    result = await graph.compile().ainvoke({"x": 0})
+    compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile()
+    initial_state: UpdateElisionState = _initial_state()
+    result: UpdateElisionState = await compiled.ainvoke(initial_state)
     assert result["x"] == 1
+    assert result["dc"] == DataClassPayload(1)
+    assert result["model"].value == 1
+    assert result["td"] == {"flag": True, "n": 1}
+    assert result["obj"] == {"n": 1}
+    assert result["items"] == [1]
+
+
+async def test_changed_slow_update_overrides_fast_update() -> None:
+    graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
+
+    async def start_node(state: UpdateElisionState) -> Command:
+        return Command(goto=[Send("fast_node", None), Send("slow_node", None)])
+
+    async def fast_node(state: UpdateElisionState) -> UpdateElisionState:
+        state.x = 1
+        state.dc.value = 1
+        state.model.value = 1
+        state.td["flag"] = True
+        state.td["n"] = 1
+        state.obj["n"] = 1
+        state.items.append(1)
+        return state
+
+    async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
+        await asyncio.sleep(0.1)
+        # Slow node makes real changes for all field types.
+        state.x = 2
+        state.dc.value = 2
+        state.model.value = 2
+        state.td["flag"] = False
+        state.td["n"] = 2
+        state.obj["n"] = 2
+        state.items.append(2)
+        return state
+
+    graph.add_entry_node(start_node)
+    graph.add_node(fast_node)
+    graph.add_finish_node(slow_node)
+
+    compiled: CompiledGraphEngine[UpdateElisionState] = graph.compile()
+    initial_state: UpdateElisionState = _initial_state()
+    result: UpdateElisionState = await compiled.ainvoke(initial_state)
+    assert result["x"] == 2
+    assert result["dc"] == DataClassPayload(2)
+    assert result["model"].value == 2
+    assert result["td"] == {"flag": False, "n": 2}
+    assert result["obj"] == {"n": 2}
+    assert result["items"] == [2]