diff --git a/packages/vite/src/plugins/__tests__/vite-plugin-rsc-transform-server.test.ts b/packages/vite/src/plugins/__tests__/vite-plugin-rsc-transform-server.test.ts
new file mode 100644
index 000000000000..b42f5dd0b02e
--- /dev/null
+++ b/packages/vite/src/plugins/__tests__/vite-plugin-rsc-transform-server.test.ts
@@ -0,0 +1,285 @@
+import { vol } from 'memfs'
+import type { TransformPluginContext } from 'rollup'
+import {
+  afterAll,
+  beforeAll,
+  describe,
+  it,
+  expect,
+  vi,
+  afterEach,
+} from 'vitest'
+
+import { rscTransformUseServerPlugin } from '../vite-plugin-rsc-transform-server.js'
+
+vi.mock('fs', async () => ({ default: (await import('memfs')).fs }))
+
+const RWJS_CWD = process.env.RWJS_CWD
+
+beforeAll(() => {
+  process.env.RWJS_CWD = '/Users/tobbe/rw-app/'
+
+  // Add a toml entry for getPaths et al.
+  vol.fromJSON({ 'redwood.toml': '' }, process.env.RWJS_CWD)
+})
+
+afterAll(() => {
+  process.env.RWJS_CWD = RWJS_CWD
+})
+
+function getPluginTransform() {
+  const plugin = rscTransformUseServerPlugin()
+
+  if (typeof plugin.transform !== 'function') {
+    throw new Error('Plugin does not have a transform function')
+  }
+
+  // Calling `bind` to please TS
+  // See https://stackoverflow.com/a/70463512/88106
+  // Typecasting because we're only going to call transform, and we don't need
+  // anything provided by the context.
+  return plugin.transform.bind({} as TransformPluginContext)
+}
+
+const pluginTransform = getPluginTransform()
+
+describe('rscTransformUseServerPlugin module scoped "use server"', () => {
+  afterEach(() => {
+    vi.resetAllMocks()
+  })
+
+  it('should handle one function', async () => {
+    const id = 'some/path/to/actions.ts'
+    const input = `
+      'use server'
+
+      import fs from 'node:fs'
+
+      export async function formAction(formData: FormData) {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\`
+        )
+      }`.trim()
+
+    const output = await pluginTransform(input, id)
+
+    expect(output).toMatchInlineSnapshot(`
+      "'use server'
+
+            import fs from 'node:fs'
+
+            export async function formAction(formData: FormData) {
+              await fs.promises.writeFile(
+                'settings.json',
+                \`{ "delay": \${formData.get('delay')} }\`
+              )
+            }
+
+      import {registerServerReference} from "react-server-dom-webpack/server";
+      registerServerReference(formAction,"some/path/to/actions.ts","formAction");
+      "
+    `)
+  })
+
+  it('should handle two functions', async () => {
+    const id = 'some/path/to/actions.ts'
+    const input = `
+      'use server'
+
+      import fs from 'node:fs'
+
+      export async function formAction1(formData: FormData) {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\`
+        )
+      }
+
+      export async function formAction2(formData: FormData) {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\`
+        )
+      }`.trim()
+
+    const output = await pluginTransform(input, id)
+
+    expect(output).toMatchInlineSnapshot(`
+      "'use server'
+
+            import fs from 'node:fs'
+
+            export async function formAction1(formData: FormData) {
+              await fs.promises.writeFile(
+                'settings.json',
+                \`{ "delay": \${formData.get('delay')} }\`
+              )
+            }
+
+            export async function formAction2(formData: FormData) {
+              await fs.promises.writeFile(
+                'settings.json',
+                \`{ "delay": \${formData.get('delay')} }\`
+              )
+            }
+
+      import {registerServerReference} from "react-server-dom-webpack/server";
+      registerServerReference(formAction1,"some/path/to/actions.ts","formAction1");
+      registerServerReference(formAction2,"some/path/to/actions.ts","formAction2");
+      "
+    `)
+  })
+
+  it('should handle arrow function', async () => {
+    const id = 'some/path/to/actions.ts'
+    const input = `
+      'use server'
+
+      import fs from 'node:fs'
+
+      export const formAction = async (formData: FormData) => {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\`
+        )
+      }`.trim()
+
+    const output = await pluginTransform(input, id)
+
+    expect(output).toMatchInlineSnapshot(`
+      "'use server'
+
+            import fs from 'node:fs'
+
+            export const formAction = async (formData: FormData) => {
+              await fs.promises.writeFile(
+                'settings.json',
+                \`{ "delay": \${formData.get('delay')} }\`
+              )
+            }
+
+      import {registerServerReference} from "react-server-dom-webpack/server";
+      if (typeof formAction === "function") registerServerReference(formAction,"some/path/to/actions.ts","formAction");
+      "
+    `)
+  })
+
+  it.todo('should handle default exported arrow function', async () => {
+    const id = 'some/path/to/actions.ts'
+    const input = `
+      'use server'
+
+      import fs from 'node:fs'
+
+      export default async (formData: FormData) => {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\`
+        )
+      }`.trim()
+
+    const output = await pluginTransform(input, id)
+
+    expect(output).toMatchInlineSnapshot(`
+      "'use server'
+
+            import fs from 'node:fs'
+
+            export default async (formData: FormData) => {
+              await fs.promises.writeFile(
+                'settings.json',
+                \`{ "delay": \${formData.get('delay')} }\`
+              )
+            }
+
+      import {registerServerReference} from "react-server-dom-webpack/server";
+      registerServerReference(default,"some/path/to/actions.ts","default");
+      "
+    `)
+  })
+
+  it('should handle default exported named function', async () => {
+    const id = 'some/path/to/actions.ts'
+    const input = `
+      "use server"
+
+      import fs from 'node:fs'
+
+      export default async function formAction(formData: FormData) {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\`
+        )
+      }`.trim()
+
+    const output = await pluginTransform(input, id)
+
+    expect(output).toMatchInlineSnapshot(`
+      ""use server"
+
+            import fs from 'node:fs'
+
+            export default async function formAction(formData: FormData) {
+              await fs.promises.writeFile(
+                'settings.json',
+                \`{ "delay": \${formData.get('delay')} }\`
+              )
+            }
+
+      import {registerServerReference} from "react-server-dom-webpack/server";
+      registerServerReference(formAction,"some/path/to/actions.ts","default");
+      "
+    `)
+  })
+
+  it.todo('should handle default exported anonymous function', async () => {
+    const id = 'some/path/to/actions.ts'
+    const input = `
+      'use server'
+
+      import fs from 'node:fs'
+
+      export default async function (formData: FormData) {
+        await fs.promises.writeFile(
+          'settings.json',
+          \`{ "delay": \${formData.get('delay')} }\n\`
+        )
+      }`
+
+    const output = await pluginTransform(input, id)
+
+    if (typeof output !== 'string') {
+      throw new Error('Expected output to be a string')
+    }
+
+    // Check that the file has a "use server" directive at the top
+    // Comments and other directives are allowed before it.
+    // Maybe also imports, I'm not sure, but am going to allow it for now. If
+    // someone finds a problem with that, we can revisit.
+    const outputLines = output.split('\n')
+    const firstCodeLineIndex = outputLines.findIndex(
+      (line) =>
+        line.startsWith('export ') ||
+        line.startsWith('async ') ||
+        line.startsWith('function ') ||
+        line.startsWith('const ') ||
+        line.startsWith('let ') ||
+        line.startsWith('var '),
+    )
+    expect(
+      outputLines
+        .slice(0, firstCodeLineIndex)
+        .some((line) => line.startsWith('"use server"')),
+    ).toBeTruthy()
+    expect(output).toContain(
+      'import {registerServerReference} from "react-server-dom-webpack/server";',
+    )
+    expect(output).toContain(
+      `registerServerReference(formAction,"${id}","default");`,
+    )
+    // One import and (exactly) one call to registerServerReference, so two
+    // matches
+    expect(output.match(/registerServerReference/g)).toHaveLength(2)
+  })
+})
diff --git a/packages/vite/src/plugins/vite-plugin-rsc-transform-server.ts b/packages/vite/src/plugins/vite-plugin-rsc-transform-server.ts
index b1a1aed6d44b..9b6cc64906c0 100644
--- a/packages/vite/src/plugins/vite-plugin-rsc-transform-server.ts
+++ b/packages/vite/src/plugins/vite-plugin-rsc-transform-server.ts
@@ -30,11 +30,9 @@ export function rscTransformUseServerPlugin(): Plugin {
       let useClient = false
       let useServer = false
 
-      for (let i = 0; i < body.length; i++) {
-        const node = body[i]
-
+      for (const node of body) {
         if (node.type !== 'ExpressionStatement' || !node.directive) {
-          break
+          continue
         }
 
         if (node.directive === 'use client') {
@@ -46,17 +44,17 @@ export function rscTransformUseServerPlugin(): Plugin {
         }
       }
 
-      if (!useServer) {
-        return code
-      }
-
       if (useClient && useServer) {
         throw new Error(
           'Cannot have both "use client" and "use server" directives in the same file.',
         )
       }
 
-      const transformedCode = transformServerModule(body, id, code)
+      let transformedCode = code
+
+      if (useServer) {
+        transformedCode = transformServerModule(body, id, code)
+      }
 
       return transformedCode
     },
@@ -119,9 +117,7 @@ function transformServerModule(
   const localNames = new Map<string, string>()
   const localTypes = new Map<string, string>()
 
-  for (let i = 0; i < body.length; i++) {
-    const node = body[i]
-
+  for (const node of body) {
     switch (node.type) {
       case 'ExportAllDeclaration':
         // If export * is used, the other file needs to explicitly opt into "use server" too.
@@ -137,7 +133,7 @@ function transformServerModule(
           }
         }
 
-        continue
+        break
 
       case 'ExportNamedDeclaration':
         if (node.declaration) {
@@ -173,12 +169,11 @@ function transformServerModule(
           }
         }
 
-        continue
+        break
     }
   }
 
   let newSrc =
-    '"use server"\n' +
     code +
     '\n\n' +
     'import {registerServerReference} from ' +