diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
index fe923a5290..16f47068ee 100644
--- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
@@ -1646,7 +1646,7 @@ function codegenInstructionValue(
       ).unwrap();
       if (instrValue.expr.type === "ArrowFunctionExpression") {
         let body: t.BlockStatement | t.Expression = fn.body;
-        if (body.body.length === 1) {
+        if (body.body.length === 1 && loweredFunc.directives.length == 0) {
           const stmt = body.body[0]!;
           if (stmt.type === "ReturnStatement" && stmt.argument != null) {
             body = stmt.argument;
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.expect.md
new file mode 100644
index 0000000000..cefe04340c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.expect.md
@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+function Component() {
+  "use strict";
+  let [count, setCount] = React.useState(0);
+  const update = () => {
+    "worklet";
+    setCount((count) => count + 1);
+  };
+  return <button onClick={update}>{count}</button>;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component() {
+  "use strict";
+  const $ = useMemoCache(3);
+
+  const [count, setCount] = React.useState(0);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      "worklet";
+
+      setCount((count_0) => count_0 + 1);
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const update = t0;
+  let t1;
+  if ($[1] !== count) {
+    t1 = <button onClick={update}>{count}</button>;
+    $[1] = count;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.js
new file mode 100644
index 0000000000..0b375af64f
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-expr-directive.js
@@ -0,0 +1,9 @@
+function Component() {
+  "use strict";
+  let [count, setCount] = React.useState(0);
+  const update = () => {
+    "worklet";
+    setCount((count) => count + 1);
+  };
+  return <button onClick={update}>{count}</button>;
+}
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.expect.md
new file mode 100644
index 0000000000..cd9fba6de6
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.expect.md
@@ -0,0 +1,50 @@
+
+## Input
+
+```javascript
+function useFoo() {
+  const update = () => {
+    "worklet";
+    return 1;
+  };
+  return update;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+  isComponent: false,
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function useFoo() {
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      "worklet";
+      return 1;
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const update = t0;
+  return update;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+  isComponent: false,
+};
+
+```
+      
+### Eval output
+(kind: ok) "[[ function params=0 ]]"
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.js
new file mode 100644
index 0000000000..da2a288376
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/arrow-function-one-line-directive.js
@@ -0,0 +1,13 @@
+function useFoo() {
+  const update = () => {
+    "worklet";
+    return 1;
+  };
+  return update;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+  isComponent: false,
+};