Merge pull request #257 from PrefectHQ/fix-flow-context

Fix flow context when using @flow
PrefectHQ · Sep 2, 2024 · ddebeb8 · ddebeb8
2 parents 6a1b384 + 5c8575b
commit ddebeb8
Showing 1 changed file with 10 additions and 10 deletions.
diff --git a/src/controlflow/decorators.py b/src/controlflow/decorators.py
@@ -97,16 +97,16 @@ def wrapper(
         if agents is not None:
             flow_kwargs.setdefault("agents", agents)
 
-        flow_obj = Flow(
-            name=fn.__name__,
-            description=fn.__doc__,
-            context=bound.arguments,
-            **flow_kwargs,
-        )
-
-        with flow_obj.create_context(create_prefect_flow_context=False):
-            with controlflow.instructions(instructions):
-                return fn(*wrapper_args, **wrapper_kwargs)
+        with (
+            Flow(
+                name=fn.__name__,
+                description=fn.__doc__,
+                context=bound.arguments,
+                **flow_kwargs,
+            ),
+            controlflow.instructions(instructions),
+        ):
+            return fn(*wrapper_args, **wrapper_kwargs)
 
     return wrapper