test: add comprehensive tests for _sanitize_for_json

27 tests covering:
- Passthrough of valid types (floats, ints, strings, booleans, None)
- NaN and +/-Infinity replacement with None
- Recursive sanitization through dicts, lists, and tuples
- Real-world event structures (on_chain_end, streaming chunks, state sync)

Author: jpr5

sdk-python/tests/test_sanitize_for_json.py

@@ -0,0 +1,211 @@
+"""Tests for _sanitize_for_json in langgraph_agent.py.
+
+Covers the NaN/Infinity sanitization fix for issue #1955:
+  - Valid floats, ints, strings, booleans, None pass through unchanged
+  - NaN and +/-Infinity are replaced with None
+  - Sanitization recurses into dicts, lists, and tuples
+  - Nested structures with mixed valid/invalid values are handled correctly
+"""
+
+import math
+
+from copilotkit.langgraph_agent import _sanitize_for_json
+
+
+class TestSanitizePassthrough:
+    """Values that should pass through unchanged."""
+
+    def test_regular_float(self):
+        assert _sanitize_for_json(3.14) == 3.14
+
+    def test_zero_float(self):
+        assert _sanitize_for_json(0.0) == 0.0
+
+    def test_negative_float(self):
+        assert _sanitize_for_json(-1.5) == -1.5
+
+    def test_integer(self):
+        assert _sanitize_for_json(42) == 42
+
+    def test_string(self):
+        assert _sanitize_for_json("hello") == "hello"
+
+    def test_boolean_true(self):
+        assert _sanitize_for_json(True) is True
+
+    def test_boolean_false(self):
+        assert _sanitize_for_json(False) is False
+
+    def test_none(self):
+        assert _sanitize_for_json(None) is None
+
+    def test_empty_dict(self):
+        assert _sanitize_for_json({}) == {}
+
+    def test_empty_list(self):
+        assert _sanitize_for_json([]) == []
+
+    def test_dict_with_valid_values(self):
+        data = {"a": 1, "b": "two", "c": 3.0, "d": None, "e": True}
+        assert _sanitize_for_json(data) == data
+
+    def test_list_with_valid_values(self):
+        data = [1, "two", 3.0, None, True]
+        assert _sanitize_for_json(data) == data
+
+
+class TestSanitizeNanInfinity:
+    """NaN and Infinity values must be replaced with None."""
+
+    def test_nan_becomes_none(self):
+        assert _sanitize_for_json(float("nan")) is None
+
+    def test_positive_infinity_becomes_none(self):
+        assert _sanitize_for_json(float("inf")) is None
+
+    def test_negative_infinity_becomes_none(self):
+        assert _sanitize_for_json(float("-inf")) is None
+
+    def test_math_nan_becomes_none(self):
+        assert _sanitize_for_json(math.nan) is None
+
+    def test_math_inf_becomes_none(self):
+        assert _sanitize_for_json(math.inf) is None
+
+
+class TestSanitizeNestedStructures:
+    """Sanitization must recurse into nested dicts and lists."""
+
+    def test_nan_in_dict_value(self):
+        result = _sanitize_for_json({"score": float("nan")})
+        assert result == {"score": None}
+
+    def test_inf_in_dict_value(self):
+        result = _sanitize_for_json({"score": float("inf")})
+        assert result == {"score": None}
+
+    def test_nan_in_list(self):
+        result = _sanitize_for_json([1.0, float("nan"), 3.0])
+        assert result == [1.0, None, 3.0]
+
+    def test_nan_in_tuple(self):
+        result = _sanitize_for_json((1.0, float("nan"), 3.0))
+        assert result == [1.0, None, 3.0]
+
+    def test_deeply_nested_nan(self):
+        data = {
+            "outer": {
+                "inner": {
+                    "values": [1.0, float("nan"), {"deep": float("inf")}]
+                }
+            }
+        }
+        result = _sanitize_for_json(data)
+        assert result == {
+            "outer": {
+                "inner": {
+                    "values": [1.0, None, {"deep": None}]
+                }
+            }
+        }
+
+    def test_mixed_valid_and_invalid_in_dict(self):
+        data = {
+            "valid_int": 42,
+            "valid_str": "hello",
+            "nan_val": float("nan"),
+            "inf_val": float("inf"),
+            "neg_inf_val": float("-inf"),
+            "valid_float": 3.14,
+            "nested": {"also_nan": float("nan"), "ok": True},
+        }
+        result = _sanitize_for_json(data)
+        assert result == {
+            "valid_int": 42,
+            "valid_str": "hello",
+            "nan_val": None,
+            "inf_val": None,
+            "neg_inf_val": None,
+            "valid_float": 3.14,
+            "nested": {"also_nan": None, "ok": True},
+        }
+
+
+class TestSanitizeEventLikeStructures:
+    """Simulate real LangGraph event structures to ensure they're sanitized.
+
+    These tests verify that the kind of data flowing through the event stream
+    (line 507 in langgraph_agent.py) is properly sanitized before JSON serialization.
+    Without sanitization, langchain_dumps() would raise ValueError for NaN/Infinity.
+    """
+
+    def test_event_with_nan_in_data_output(self):
+        """Simulates an on_chain_end event where output state contains NaN."""
+        event = {
+            "event": "on_chain_end",
+            "name": "agent_node",
+            "run_id": "abc-123",
+            "metadata": {},
+            "data": {
+                "output": {
+                    "score": float("nan"),
+                    "result": "success",
+                    "confidence": float("inf"),
+                }
+            },
+        }
+        result = _sanitize_for_json(event)
+        assert result["data"]["output"]["score"] is None
+        assert result["data"]["output"]["confidence"] is None
+        assert result["data"]["output"]["result"] == "success"
+        assert result["event"] == "on_chain_end"
+
+    def test_event_with_nan_in_streaming_chunk(self):
+        """Simulates a streaming event where a chunk contains NaN values."""
+        event = {
+            "event": "on_chat_model_stream",
+            "name": "model",
+            "run_id": "def-456",
+            "metadata": {"copilotkit:emit-intermediate-state": True},
+            "data": {
+                "chunk": {
+                    "content": "",
+                    "tool_call_chunks": [
+                        {
+                            "args": '{"temperature": NaN}',
+                            "score": float("nan"),
+                        }
+                    ],
+                }
+            },
+        }
+        result = _sanitize_for_json(event)
+        assert result["data"]["chunk"]["tool_call_chunks"][0]["score"] is None
+        # String "NaN" inside args is NOT a float, so it stays as-is
+        assert result["data"]["chunk"]["tool_call_chunks"][0]["args"] == '{"temperature": NaN}'
+
+    def test_event_with_no_problematic_values_unchanged(self):
+        """Normal events with no NaN/Infinity should pass through with same structure."""
+        event = {
+            "event": "on_chain_start",
+            "name": "agent_node",
+            "run_id": "ghi-789",
+            "metadata": {"key": "value"},
+            "data": {"input": {"query": "hello", "count": 5}},
+        }
+        result = _sanitize_for_json(event)
+        assert result == event
+
+    def test_state_sync_payload_with_nan(self):
+        """Simulates the state dict passed to _emit_state_sync_event."""
+        state = {
+            "messages": [],
+            "score": float("nan"),
+            "embedding": [0.1, float("inf"), 0.3, float("-inf")],
+            "metadata": {"loss": float("nan"), "accuracy": 0.95},
+        }
+        result = _sanitize_for_json(state)
+        assert result["score"] is None
+        assert result["embedding"] == [0.1, None, 0.3, None]
+        assert result["metadata"]["loss"] is None
+        assert result["metadata"]["accuracy"] == 0.95