diff --git a/.cursor/rules.md b/.cursor/rules.md
new file mode 100644
index 0000000000..00385e6f5f
--- /dev/null
+++ b/.cursor/rules.md
@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+
diff --git a/.gemini/rules.md b/.gemini/rules.md
new file mode 100644
index 0000000000..00385e6f5f
--- /dev/null
+++ b/.gemini/rules.md
@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+
diff --git a/AGENTS.md b/AGENTS.md
new file mode 100644
index 0000000000..00385e6f5f
--- /dev/null
+++ b/AGENTS.md
@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+
diff --git a/CLAUDE.md b/CLAUDE.md
new file mode 100644
index 0000000000..00385e6f5f
--- /dev/null
+++ b/CLAUDE.md
@@ -0,0 +1,5 @@
+The AI agents should never link to any issue or a pull request
+in any GitHub repository in the code reviews!
+
+The AI agents should not review AI agents' config files like CLAUDE.md or AGENTS.md!
+
diff --git a/python/Cargo.lock b/python/Cargo.lock
index 8ccc4358b7..dc1c986dc8 100644
--- a/python/Cargo.lock
+++ b/python/Cargo.lock
@@ -642,6 +642,7 @@ dependencies = [
  "futures",
  "http",
  "insta",
+ "itertools",
  "log",
  "object_store",
  "parking_lot",
@@ -652,6 +653,7 @@ dependencies = [
  "tokio",
  "tokio-stream",
  "tonic",
+ "tower-http",
  "uuid",
 ]
 
diff --git a/python/python/ballista/extension.py b/python/python/ballista/extension.py
index 90a3651bc6..5dbbf2a85e 100644
--- a/python/python/ballista/extension.py
+++ b/python/python/ballista/extension.py
@@ -72,12 +72,11 @@ def method_wrapper(*args, **kwargs):
 
         for base_name, base_value in bases[0].__dict__.items():
             #
-            # TODO: could we not use 'DataFrame' as a string here?
             #
             if (
                 callable(base_value)
                 and not base_name.startswith("__")
-                and base_value.__annotations__.get("return") == "DataFrame"
+                and base_value.__annotations__.get("return") == DataFrame.__name__
             ):
                 #
                 # functions which return DataFrame are redefined
@@ -112,12 +111,11 @@ def method_wrapper(*args, **kwargs):
 
         for base_name, base_value in bases[0].__dict__.items():
             #
-            # could we not use 'DataFrame' as a string here?
             #
             if (
                 callable(base_value)
                 and not base_name.startswith("__")
-                and base_value.__annotations__.get("return") == "DataFrame"
+                and base_value.__annotations__.get("return") == DataFrame.__name__
             ):
                 #
                 # functions which return DataFrame are redefined
diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py
index bb92eb24c5..047c8986e7 100644
--- a/python/python/tests/test_context.py
+++ b/python/python/tests/test_context.py
@@ -16,6 +16,7 @@
 # under the License.
 
 from ballista import BallistaSessionContext, setup_test_cluster
+from ballista.extension import DataFrame, DistributedDataFrame, SessionContext
 from datafusion import col, lit
 import pytest
 import pyarrow as pa
@@ -138,3 +139,31 @@ def test_write_json(ctx, tmp_path):
     df.write_json(out_dir)
     json_files = list((tmp_path / "out").glob("*.json"))
     assert len(json_files) > 0
+
+
+def _assert_dataframe_returning_methods_wrapped(base_cls, sub_cls):
+    should_be_wrapped = {
+        name
+        for name, val in base_cls.__dict__.items()
+        if callable(val)
+        and not name.startswith("__")
+        and val.__annotations__.get("return") == DataFrame.__name__
+    }
+
+    assert should_be_wrapped
+    for name in should_be_wrapped:
+        assert name in sub_cls.__dict__, f"{name} not found in {sub_cls.__name__}"
+        assert callable(sub_cls.__dict__[name]), (
+            f"{name} is not callable in {sub_cls.__name__}"
+        )
+        assert sub_cls.__dict__[name] is not base_cls.__dict__[name], (
+            f"{name} was not replaced in {sub_cls.__name__}"
+        )
+
+
+def test_distributed_dataframe_wraps_dataframe_returning_methods():
+    _assert_dataframe_returning_methods_wrapped(DataFrame, DistributedDataFrame)
+
+
+def test_ballista_session_context_wraps_dataframe_returning_methods():
+    _assert_dataframe_returning_methods_wrapped(SessionContext, BallistaSessionContext)