Resolve Monty persistence review notes

Author: CyrusNuevoDia

rlm/core/rlm.py

@@ -363,7 +363,7 @@ def _validate_persistent_environment_support(self) -> None:
         - add_context(payload, index): Add new context for multi-turn conversations
         - get_context_count(): Return the number of loaded contexts
 
-        Currently only 'local' (LocalREPL) supports these methods.
+        Currently only 'local' (LocalREPL) and 'monty' (MontyREPL) support these methods.
 
         Raises:
             ValueError: If the environment type does not support persistent mode.

rlm/environments/monty_repl.py

@@ -175,10 +175,6 @@ def __init__(
         resource_limits: pydantic_monty.ResourceLimits | None = None,
         **kwargs: Any,
     ) -> None:
-        if persistent:
-            raise NotImplementedError(
-                "Persistent REPLs are currently not supported for environment: MontyREPL"
-            )
         super().__init__(persistent=persistent, depth=depth, **kwargs)
 
         self.lm_handler_address = lm_handler_address

tests/test_multi_turn_integration.py

@@ -348,13 +348,29 @@ def test_local_environment_supported(self):
     def test_monty_environment_supported(self):
         """Monty environment should support persistent mode."""
         pytest.importorskip("pydantic_monty")
-        rlm = RLM(
-            backend="openai",
-            backend_kwargs={"model_name": "test"},
-            environment="monty",
-            persistent=True,
-        )
-        assert rlm.persistent is True
+        responses = [
+            "```repl\nx = 2\nprint(x)\n```\nFINAL(done)",
+            "```repl\nprint(x)\n```\nFINAL(done)",
+        ]
+
+        with patch.object(rlm_module, "get_client") as mock_get_client:
+            mock_lm = create_mock_lm(responses)
+            mock_get_client.return_value = mock_lm
+
+            with RLM(
+                backend="openai",
+                backend_kwargs={"model_name": "test"},
+                environment="monty",
+                persistent=True,
+            ) as rlm:
+                rlm.completion("First turn")
+                mock_lm.completion.side_effect = list(responses[1:])
+                rlm.completion("Second turn")
+
+                env = rlm._persistent_env
+                assert env is not None
+                assert env.get_context_count() == 2
+                assert "x" in env.locals
 
 
 class TestMultiTurnEndToEnd: