revert twinapi change

OpenMined · Sep 23, 2024 · 2eca5a5 · 2eca5a5
1 parent f510485
commit 2eca5a5
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/packages/syft/src/syft/service/api/api.py b/packages/syft/src/syft/service/api/api.py
@@ -568,6 +568,16 @@ def exec_code(
             # run all this code to clean up the state
             code.update_state(internal_context.state)
 
+            if isinstance(code, PublicAPIEndpoint):
+                self.mock_function = code
+            else:
+                self.private_function = code  # type: ignore
+
+            api_service = context.server.get_service("apiservice")
+            api_service.stash.upsert(
+                context.server.services.user.root_verify_key, self
+            ).unwrap()
+
             print = original_print  # type: ignore
             # if we caught a SyftException above we will raise and auto wrap to Result
             if isinstance(result, SyftException):