diff --git a/swig/python/codegen/codegen.py b/swig/python/codegen/codegen.py
index b50947574..3f2f3b938 100644
--- a/swig/python/codegen/codegen.py
+++ b/swig/python/codegen/codegen.py
@@ -720,13 +720,14 @@ def _Assign(self, t):
         # check if target exists in locals
         if t.targets[0].id not in self._locals :
             # Special case, catch message.at() where a message is returned outside a message loop
-            if hasattr(t.value, "func") and isinstance(t.value.func, ast.Attribute) and t.value.func.attr == 'at' :
-                if t.value.func.value.id == self._input_message_var :
-                    self._standalone_message_var.append(t.targets[0].id)
-            # Special case, track which variables hold directed graph handles
-            elif hasattr(t.value, "func") and t.value.func.attr == 'getDirectedGraph' :
-                if t.value.func.value.value.id == "pyflamegpu" and t.value.func.value.attr == "environment" :
-                    self._directed_graph_vars.append(t.targets[0].id)
+            if hasattr(t.value, "func") and isinstance(t.value.func, ast.Attribute):
+                if t.value.func.attr == 'at' :
+                    if t.value.func.value.id == self._input_message_var :
+                        self._standalone_message_var.append(t.targets[0].id)
+                # Special case, track which variables hold directed graph handles
+                elif t.value.func.attr == 'getDirectedGraph' :
+                    if t.value.func.value.value.id == "pyflamegpu" and t.value.func.value.attr == "environment" :
+                        self._directed_graph_vars.append(t.targets[0].id)
             # Special case, definitions outside of agent fn are made const
             if self._indent == 0:
                 self.write("constexpr ")
diff --git a/tests/python/codegen/test_codegen.py b/tests/python/codegen/test_codegen.py
index 18f810c89..6f51a08f6 100644
--- a/tests/python/codegen/test_codegen.py
+++ b/tests/python/codegen/test_codegen.py
@@ -382,72 +382,84 @@ def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone
 """
 
 py_fgpu_graph_fns = """\
-fgraph = pyflamegpu.environment.getDirectedGraph("fgraph")
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    fgraph = pyflamegpu.environment.getDirectedGraph("fgraph")
 
-# Fetch the ID of the vertex at index 0
-vertex_id = fgraph.getVertexID(0)
-# Fetch the index of the vertex with ID 1
-vertex_index = fgraph.getVertexIndex(1)
+    # Fetch the ID of the vertex at index 0
+    vertex_id = fgraph.getVertexID(0)
+    # Fetch the index of the vertex with ID 1
+    vertex_index = fgraph.getVertexIndex(1)
 
-# Access a property of vertex with ID 1
-bar_0 = fgraph.getVertexPropertyFloatArray2("bar", 0)
+    # Access a property of vertex with ID 1
+    bar_0 = fgraph.getVertexPropertyFloatArray2("bar", 0)
 
-# Fetch the source and destination indexes from the edge at index 0
-source_index = fgraph.getEdgeSource(0)
-destination_index = fgraph.getEdgeDestination(0)
+    # Fetch the source and destination indexes from the edge at index 0
+    source_index = fgraph.getEdgeSource(0)
+    destination_index = fgraph.getEdgeDestination(0)
 
-# Fetch the index of the edge from vertex ID 1 to vertex ID 2
-edge_index = fgraph.getEdgeIndex(1, 2)
+    # Fetch the index of the edge from vertex ID 1 to vertex ID 2
+    edge_index = fgraph.getEdgeIndex(1, 2)
 
-# Access a property of edge with source ID 1, destination ID 2
-foo2 = fgraph.getEdgePropertyInt("foo", edge_index);
+    # Access a property of edge with source ID 1, destination ID 2
+    foo2 = fgraph.getEdgePropertyInt("foo", edge_index);
 """
 cpp_fgpu_graph_fns = """\
-auto fgraph = FLAMEGPU->environment.getDirectedGraph("fgraph");
-auto vertex_id = fgraph.getVertexID(0);
-auto vertex_index = fgraph.getVertexIndex(1);
-auto bar_0 = fgraph.getVertexProperty<float, 2>("bar", 0);
-auto source_index = fgraph.getEdgeSource(0);
-auto destination_index = fgraph.getEdgeDestination(0);
-auto edge_index = fgraph.getEdgeIndex(1, 2);
-auto foo2 = fgraph.getEdgeProperty<int>("foo", edge_index);
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto fgraph = FLAMEGPU->environment.getDirectedGraph("fgraph");
+    auto vertex_id = fgraph.getVertexID(0);
+    auto vertex_index = fgraph.getVertexIndex(1);
+    auto bar_0 = fgraph.getVertexProperty<float, 2>("bar", 0);
+    auto source_index = fgraph.getEdgeSource(0);
+    auto destination_index = fgraph.getEdgeDestination(0);
+    auto edge_index = fgraph.getEdgeIndex(1, 2);
+    auto foo2 = fgraph.getEdgeProperty<int>("foo", edge_index);
+}
 """
 
 py_fgpu_for_graph_in_fns = """\
-# Iterate the edges joining the vertex with ID 1
-fgraph = pyflamegpu.environment.getDirectedGraph("fgraph")
-for edge in fgraph.inEdges(vertex_index):
-    # Read the current edges' source vertex index
-    src_vertex_index = edge.getEdgeSource()
-    # Read a property from the edge
-    foo = edge.getPropertyInt("foo")
-    bar = edge.getPropertyFloatArray2("bar", 0)
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    # Iterate the edges joining the vertex with ID 1
+    fgraph = pyflamegpu.environment.getDirectedGraph("fgraph")
+    for edge in fgraph.inEdges(vertex_index):
+        # Read the current edges' source vertex index
+        src_vertex_index = edge.getEdgeSource()
+        # Read a property from the edge
+        foo = edge.getPropertyInt("foo")
+        bar = edge.getPropertyFloatArray2("bar", 0)
 """
 cpp_fgpu_for_graph_in_fns = """\
-auto fgraph = FLAMEGPU->environment.getDirectedGraph("fgraph");
-for (const auto& edge : fgraph.inEdges(vertex_index)){
-    auto src_vertex_index = edge.getEdgeSource();
-    auto foo = edge.getProperty<int>("foo");
-    auto bar = edge.getProperty<float, 2>("bar", 0);
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto fgraph = FLAMEGPU->environment.getDirectedGraph("fgraph");
+    for (const auto& edge : fgraph.inEdges(vertex_index)){
+        auto src_vertex_index = edge.getEdgeSource();
+        auto foo = edge.getProperty<int>("foo");
+        auto bar = edge.getProperty<float, 2>("bar", 0);
+    }
 }
 """
 
 py_fgpu_for_graph_out_fns = """\
-# Iterate the edges leaving the vertex with ID 1
-fgraph = pyflamegpu.environment.getDirectedGraph("fgraph")
-for edge in fgraph.outEdges(vertex_index):
-    # Read the current edges' destination vertex index
-    dest_vertex_index = edge.getEdgeDestination()
-    # Read a property from the edge
-    foo = edge.getPropertyInt("foo")
-    bar = edge.getPropertyFloatArray2("bar", 0)
+@pyflamegpu.agent_function
+def func(message_in: pyflamegpu.MessageNone, message_out: pyflamegpu.MessageNone) :
+    # Iterate the edges leaving the vertex with ID 1
+    fgraph = pyflamegpu.environment.getDirectedGraph("fgraph")
+    for edge in fgraph.outEdges(vertex_index):
+        # Read the current edges' destination vertex index
+        dest_vertex_index = edge.getEdgeDestination()
+        # Read a property from the edge
+        foo = edge.getPropertyInt("foo")
+        bar = edge.getPropertyFloatArray2("bar", 0)
 """
 cpp_fgpu_for_graph_out_fns = """\
-auto fgraph = FLAMEGPU->environment.getDirectedGraph("fgraph");
-for (const auto& edge : fgraph.outEdges(vertex_index)){
-    auto dest_vertex_index = edge.getEdgeDestination();
-    auto foo = edge.getProperty<int>("foo");
-    auto bar = edge.getProperty<float, 2>("bar", 0);
+FLAMEGPU_AGENT_FUNCTION(func, flamegpu::MessageNone, flamegpu::MessageNone){
+    auto fgraph = FLAMEGPU->environment.getDirectedGraph("fgraph");
+    for (const auto& edge : fgraph.outEdges(vertex_index)){
+        auto dest_vertex_index = edge.getEdgeDestination();
+        auto foo = edge.getProperty<int>("foo");
+        auto bar = edge.getProperty<float, 2>("bar", 0);
+    }
 }
 """
 
diff --git a/tests/test_rtcfunc_file b/tests/test_rtcfunc_file
new file mode 100644
index 000000000..5a0166c8d
--- /dev/null
+++ b/tests/test_rtcfunc_file
@@ -0,0 +1,5 @@
+
+FLAMEGPU_AGENT_FUNCTION(rtc_test_filefunc, flamegpu::MessageNone, flamegpu::MessageNone) {
+    FLAMEGPU->setVariable<int>("x", FLAMEGPU->getVariable<int>("x") + 1);
+    return flamegpu::ALIVE;
+}