From 558e690a7fe2d90b9cd64259f96c32d9775cf2c6 Mon Sep 17 00:00:00 2001
From: Joaquin Coromina <bjoaquinc@gmail.com>
Date: Sat, 6 Sep 2025 00:36:25 +0200
Subject: [PATCH] Fix pydantic parsing error of VariableValue

---
 marimo/_mcp/server/tools/cells.py     | 17 ++++++++++++++---
 tests/_mcp/server/tools/test_cells.py | 16 ++++++++++++++--
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/marimo/_mcp/server/tools/cells.py b/marimo/_mcp/server/tools/cells.py
index 305c6a52945..cfe4d000682 100644
--- a/marimo/_mcp/server/tools/cells.py
+++ b/marimo/_mcp/server/tools/cells.py
@@ -11,7 +11,6 @@
 from marimo._mcp.server.responses import (
     SuccessResult,
 )
-from marimo._messaging.ops import VariableValue
 from marimo._server.api.deps import AppStateBase
 from marimo._types.ids import CellId_t, SessionId
 
@@ -69,7 +68,14 @@ class CellRuntimeMetadata:
     execution_time: Optional[float]
 
 
-CellVariables = dict[str, VariableValue]
+@dataclass(kw_only=True)
+class CellVariableValue:
+    name: str
+    value: Optional[str] = None
+    datatype: Optional[str] = None
+
+
+CellVariables = dict[str, CellVariableValue]
 
 
 @dataclass(kw_only=True)
@@ -407,6 +413,11 @@ def _get_cell_variables(
     for var_name in cell_defs:
         if var_name in all_variables:
             var_value = all_variables[var_name]
-            cell_variables[var_name] = var_value
+            # Convert VariableValue to pydantic parsable dataclass
+            cell_variables[var_name] = CellVariableValue(
+                name=var_name,
+                value=var_value.value,
+                datatype=var_value.datatype,
+            )
 
     return cell_variables
diff --git a/tests/_mcp/server/tools/test_cells.py b/tests/_mcp/server/tools/test_cells.py
index ffe61daffc9..772ef48a7ef 100644
--- a/tests/_mcp/server/tools/test_cells.py
+++ b/tests/_mcp/server/tools/test_cells.py
@@ -12,6 +12,7 @@
     CellErrors,
     CellRuntimeMetadata,
     CellVariables,
+    CellVariableValue,
     SupportedCellType,
     _determine_cell_type,
     _get_cell_errors,
@@ -241,7 +242,14 @@ def test_get_cell_variables_with_variables():
 
     result = _get_cell_variables(session, cell_data)
 
-    expected: CellVariables = {"x": var_x, "y": var_y}
+    expected: CellVariables = {
+        "x": CellVariableValue(
+            name="x", value=var_x.value, datatype=var_x.datatype
+        ),
+        "y": CellVariableValue(
+            name="y", value=var_y.value, datatype=var_y.datatype
+        ),
+    }
     assert result == expected
     assert (
         "z" not in result
@@ -262,5 +270,9 @@ def test_get_cell_variables_missing_variables():
     result = _get_cell_variables(session, cell_data)
 
     # Should only include variables that exist in session
-    expected: CellVariables = {"x": var_x}
+    expected: CellVariables = {
+        "x": CellVariableValue(
+            name="x", value=var_x.value, datatype=var_x.datatype
+        )
+    }
     assert result == expected