From 5c8575bdbac6a0d6edab79e539743c1cfc564975 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Sun, 1 Sep 2024 20:51:53 -0400
Subject: [PATCH] Fix flow context when using @flow

---
 src/controlflow/decorators.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/controlflow/decorators.py b/src/controlflow/decorators.py
index f605eee2..d143003a 100644
--- 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