Ensure we correctly do effect validation of tools

Author: benh

reboot/mcp/server.py

@@ -24,20 +24,21 @@
 from mcp.shared.message import ServerMessageMetadata
 from rbt.mcp.v1.session_rbt import Session
 from reboot.aio.applications import Application
-from reboot.aio.contexts import WorkflowContext
+from reboot.aio.contexts import EffectValidation, WorkflowContext
 from reboot.aio.external import ExternalContext
 from reboot.aio.types import StateRef
+from reboot.aio.workflows import at_least_once
 from reboot.mcp.event_store import DurableEventStore, replay
 from reboot.mcp.servicers.session import (
     SessionServicer,
     _servers,
     _context,
 )
-from reboot.aio.workflows import at_least_once
 from reboot.mcp.servicers.stream import StreamServicer
 from reboot.std.collections.v1 import sorted_map
 from rebootdev.aio.headers import CONSENSUS_ID_HEADER, STATE_REF_HEADER
 from rebootdev.memoize.v1.memoize_rbt import Memoize
+from rebootdev.settings import DOCS_BASE_URL
 from starlette.applications import Starlette
 from starlette.requests import Request
 from starlette.responses import Response, StreamingResponse
@@ -663,12 +664,12 @@ def _wrap_tool(fn: mcp.types.AnyFunction) -> mcp.types.AnyFunction:
 
     wrapper_signature = signature.replace(parameters=wrapper_parameters)
 
-    async def wrapper(ctx: fastmcp.Context, *args, **kwargs):
-
-        context: WorkflowContext | None = _context.get()
-
-        assert context is not None
-
+    async def wrapper(
+        ctx: fastmcp.Context,
+        context: WorkflowContext,
+        *args,
+        **kwargs,
+    ):
         # To account for the lack of "intersection" types in
         # Python (which is actively being worked on), we instead
         # create a new dynamic `DurableContext` instance that
@@ -938,19 +939,49 @@ async def send_request_and_wait_for_result():
 
             return fn(**dict(bound.arguments))
         except:
-            # TODO: print stack trace after we've fixed `memoize`
-            # effect validation bug.
-            #
-            # import traceback
-            # traceback.print_exc()
+            import traceback
+            traceback.print_exc()
             raise
 
-    setattr(wrapper, "__signature__", wrapper_signature)
-    wrapper.__name__ = fn.__name__
-    wrapper.__doc__ = fn.__doc__
+    async def wrapper_validating_effects(
+        ctx: fastmcp.Context,
+        *args,
+        **kwargs,
+    ):
+        context: WorkflowContext | None = _context.get()
+
+        assert context is not None
+
+        # Checkpoint the context since it is the `IdempotencyManager`.
+        checkpoint = context.checkpoint()
+
+        result = await wrapper(ctx, context, *args, **kwargs)
+
+        if context._effect_validation == EffectValidation.DISABLED:
+            return result
+
+        # Effect validation is enabled.
+        logger.info(
+            f"Re-running tool '{fn.__name__}' "
+            f"to validate effects. See {DOCS_BASE_URL}/develop/side_effects "
+            "for more information."
+        )
+
+        # Restore the context to the checkpoint we took above so we
+        # can re-execute `callable` as though it is being retried from
+        # scratch.
+        context.restore(checkpoint)
+
+        # TODO: check if `result` is different (we don't do this for
+        # other effect validation so we're also not doing it now).
+
+        return await wrapper(ctx, context, *args, **kwargs)
 
-    return wrapper
+    setattr(wrapper_validating_effects, "__signature__", wrapper_signature)
+    wrapper_validating_effects.__name__ = fn.__name__
+    wrapper_validating_effects.__doc__ = fn.__doc__
 
+    return wrapper_validating_effects
 
 
 class StreamableHTTPASGIApp:

reboot/mcp/servicers/session.py

@@ -1,7 +1,7 @@
+import anyio
 import mcp.types
 import pickle
 from contextvars import ContextVar
-from anyio import create_memory_object_stream
 from anyio.streams.memory import (
     MemoryObjectReceiveStream,
     MemoryObjectSendStream,
@@ -76,9 +76,10 @@ def _get_request_streams(self, request_id: mcp.types.RequestId):
                 # Create streams for communicating with MCP server.
                 self._request_streams[request_id] = Streams(
                     refs=1,  # Initial reference count.
-                    read_stream=create_memory_object_stream[
+                    read_stream=anyio.create_memory_object_stream[
                         SessionMessage | Exception](),
-                    write_stream=create_memory_object_stream[SessionMessage](),
+                    write_stream=anyio.create_memory_object_stream[
+                        SessionMessage](),
                 )
             else:
                 self._request_streams[request_id].refs += 1
@@ -121,8 +122,13 @@ async def HandleMessage(
                 )
 
                 async def send_and_receive():
-
-                    await read_stream_send.send(message)
+                    try:
+                        await read_stream_send.send(message)
+                    except anyio.ClosedResourceError:
+                        # Stream is closed, we must be re-executing
+                        # this function due to effect validation, just
+                        # return.
+                        return
 
                     async for write_message in write_stream_receive:
                         logger.debug(
@@ -141,7 +147,9 @@ async def send_and_receive():
                             write_message.message.root.id = event_id
                             related_request_id = write_message.metadata.related_request_id
                             assert related_request_id is not None
-                            self._write_request_ids[event_id] = (write_request_id, related_request_id)
+                            self._write_request_ids[event_id] = (
+                                write_request_id, related_request_id
+                            )
 
                         await stream.per_workflow(event_id).Put(
                             context,
@@ -153,6 +161,7 @@ async def send_and_receive():
                             write_message.message.root,
                             mcp.types.JSONRPCResponse | mcp.types.JSONRPCError,
                         ):
+                            await read_stream_send.aclose()
                             break
 
                 await at_least_once(
@@ -161,8 +170,10 @@ async def send_and_receive():
                     send_and_receive,
                 )
 
-                await read_stream_send.aclose()
-
+                # NOTE: need to await `run_task` within the
+                # `self._get_request_streams()` context manager so
+                # that we continue to use the same streams between
+                # this function and `Run()`.
                 await run_task
 
                 logger.debug(f"Completed ({type(message).__name__}): {message}")
@@ -221,10 +232,11 @@ async def Run(
             write_stream_send, _ = write_stream
 
             async def server_run():
-                global _servers
-                server = _servers[path]
+                assert _context.get() is None
                 _context.set(context)
                 try:
+                    global _servers
+                    server = _servers[path]
                     await server.run(
                         read_stream_receive,
                         write_stream_send,