diff --git a/jest.config.js b/jest.config.js
index e3766adb..dff8da68 100644
--- a/jest.config.js
+++ b/jest.config.js
@@ -9,7 +9,5 @@ module.exports = {
   moduleNameMapper: {
     '@nuxtjs/composition-api/dist/runtime/globals':
       '<rootDir>/src/runtime/globals',
-    '@nuxtjs/composition-api/dist/runtime/register':
-      '<rootDir>/src/runtime/register',
   },
 }
diff --git a/package.json b/package.json
index 4aec6fa2..1c407b9d 100644
--- a/package.json
+++ b/package.json
@@ -23,14 +23,15 @@
       "import": "./dist/runtime/index.mjs",
       "require": "./dist/runtime/index.js"
     },
-    "./module": "./dist/module.js",
+    "./module": "./dist/module/index.js",
     "./package.json": "./package.json",
-    "./dist/babel-plugin": "./dist/babel-plugin.js",
-    "./dist/runtime/*": "./dist/runtime/*"
+    "./dist/babel-plugin": "./dist/babel-plugin/index.js",
+    "./dist/runtime/globals": "./dist/runtime/globals.js",
+    "./dist/runtime/templates/*": "./dist/runtime/templates/*"
   },
   "main": "./dist/runtime/index.js",
   "module": "./dist/runtime/index.mjs",
-  "types": "./dist/index.d.ts",
+  "types": "./dist/runtime/index.d.ts",
   "files": [
     "dist",
     "module.js"
diff --git a/siroc.config.js b/siroc.config.js
index 90db12ef..724b6a30 100644
--- a/siroc.config.js
+++ b/siroc.config.js
@@ -2,9 +2,6 @@ import { defineSirocConfig } from 'siroc'
 
 export default defineSirocConfig({
   rollup: {
-    externals: [
-      '@nuxtjs/composition-api/dist/runtime/globals',
-      '@nuxtjs/composition-api/dist/runtime/register',
-    ],
+    externals: ['@nuxtjs/composition-api/dist/runtime/globals'],
   },
 })
diff --git a/src/babel-plugin.js b/src/babel-plugin.js
deleted file mode 100644
index c4be3fc3..00000000
--- a/src/babel-plugin.js
+++ /dev/null
@@ -1,2 +0,0 @@
-// eslint-disable-next-line
-module.exports = require('jiti')(__dirname)('./babel-plugin.ts')
diff --git a/src/babel-plugin/index.js b/src/babel-plugin/index.js
new file mode 100644
index 00000000..c8116081
--- /dev/null
+++ b/src/babel-plugin/index.js
@@ -0,0 +1,2 @@
+// eslint-disable-next-line
+module.exports = require('jiti')(__dirname)('./index.ts')
diff --git a/src/babel-plugin.ts b/src/babel-plugin/index.ts
similarity index 100%
rename from src/babel-plugin.ts
rename to src/babel-plugin/index.ts
diff --git a/src/index.ts b/src/index.ts
deleted file mode 100644
index aa642c2f..00000000
--- a/src/index.ts
+++ /dev/null
@@ -1,2 +0,0 @@
-/* dummy entry for types field */
-export * from './runtime'
diff --git a/src/babel-register.ts b/src/module/babel-register.ts
similarity index 94%
rename from src/babel-register.ts
rename to src/module/babel-register.ts
index 4cd893d1..615a1db1 100644
--- a/src/babel-register.ts
+++ b/src/module/babel-register.ts
@@ -17,7 +17,7 @@ export function registerBabelPlugin(this: ModuleThis) {
       'Unable to automatically add Babel plugin. Make sure your custom `build.babel.plugins` returns `@nuxtjs/composition-api/dist/babel-plugin`'
     )
   } else {
-    nuxtOptions.build.babel.plugins.push(resolveRelativePath('babel-plugin'))
+    nuxtOptions.build.babel.plugins.push(resolveRelativePath('../babel-plugin'))
   }
 
   /**
diff --git a/src/globals-register.ts b/src/module/globals-register.ts
similarity index 100%
rename from src/globals-register.ts
rename to src/module/globals-register.ts
diff --git a/src/module.ts b/src/module/index.ts
similarity index 81%
rename from src/module.ts
rename to src/module/index.ts
index 3c0e3bef..7757ab79 100644
--- a/src/module.ts
+++ b/src/module/index.ts
@@ -1,7 +1,8 @@
 import type { Module, NuxtOptions } from '@nuxt/types'
 import { resolve } from 'upath'
+import type { Compiler } from 'webpack'
 
-import { name, version } from '../package.json'
+import { name, version } from '../../package.json'
 
 import { registerBabelPlugin } from './babel-register'
 import { addGlobalsFile } from './globals-register'
@@ -26,6 +27,14 @@ const compositionApiModule: Module<never> = function compositionApiModule() {
     '@vue/composition-api/dist/vue-composition-api.esm.js'
   )
 
+  // Register the Vue Composition API for webpack
+
+  const registration = addResolvedTemplate.call(this, 'register.mjs')
+  this.nuxt.hook('build:compile', ({ compiler }: { compiler: Compiler }) => {
+    const entry = compiler.options.entry as Record<string, string[]>
+    entry.app.unshift(registration)
+  })
+
   // Turn off webpack4 module context for .mjs files (as it appears to have some issues)
 
   this.extendBuild(config => {
@@ -40,9 +49,11 @@ const compositionApiModule: Module<never> = function compositionApiModule() {
 
   // If we're using nuxt-vite, register vite plugin & inject configuration
 
+  const viteMiddleware = addResolvedTemplate.call(this, 'middleware.mjs')
   this.nuxt.hook('vite:extend', async (ctx: any) => {
     const { compositionApiPlugin } = await import('./vite-plugin')
     ctx.config.plugins.push(compositionApiPlugin())
+    ctx.config.resolve.alias['./middleware.js'] = viteMiddleware
   })
 
   // If we're using Babel, register Babel plugin for injecting keys
diff --git a/src/static.ts b/src/module/static.ts
similarity index 100%
rename from src/static.ts
rename to src/module/static.ts
diff --git a/src/utils.ts b/src/module/utils.ts
similarity index 89%
rename from src/utils.ts
rename to src/module/utils.ts
index 70d5783c..ed38e98c 100644
--- a/src/utils.ts
+++ b/src/module/utils.ts
@@ -1,6 +1,6 @@
 import { NuxtConfig, NuxtOptions } from '@nuxt/types'
 import { ModuleThis } from '@nuxt/types/config/module'
-import { basename, join, resolve } from 'upath'
+import { join, resolve } from 'upath'
 
 export function isFullStatic(options: NuxtConfig) {
   return (
@@ -26,10 +26,11 @@ export function addResolvedTemplate(
 ) {
   const nuxtOptions: NuxtOptions = this.nuxt.options
 
-  const src = resolveRelativePath(join('runtime/templates', template))
+  const src = resolveRelativePath(`../runtime/templates/${template}`)
+  const filename = template.replace('register.mjs', 'register.js')
   const { dst } = this.addTemplate({
     src,
-    fileName: join('composition-api', basename(src)),
+    fileName: join('composition-api', filename),
     options,
   })
 
diff --git a/src/vite-plugin.ts b/src/module/vite-plugin.ts
similarity index 100%
rename from src/vite-plugin.ts
rename to src/module/vite-plugin.ts
diff --git a/src/runtime/composables/index.ts b/src/runtime/composables/index.ts
index 1a81e4c3..18a31130 100644
--- a/src/runtime/composables/index.ts
+++ b/src/runtime/composables/index.ts
@@ -1,5 +1,3 @@
-import '../register'
-
 export { useAsync } from './async'
 export { defineComponent } from './component'
 export { useContext, withContext } from './context'
diff --git a/src/runtime/index.ts b/src/runtime/index.ts
index 8ef8d8b0..692bc41f 100644
--- a/src/runtime/index.ts
+++ b/src/runtime/index.ts
@@ -1 +1,8 @@
+import Vue from 'vue'
+import CompositionApi from '@vue/composition-api'
+
+if (process.env.NODE_ENV === 'test') {
+  Vue.use(CompositionApi)
+}
+
 export * from './composables'
diff --git a/src/runtime/templates/middleware.mjs b/src/runtime/templates/middleware.mjs
new file mode 100644
index 00000000..ce598a5a
--- /dev/null
+++ b/src/runtime/templates/middleware.mjs
@@ -0,0 +1,4 @@
+import './register.js'
+import middleware from '../middleware.js'
+
+export default middleware
diff --git a/src/runtime/register.ts b/src/runtime/templates/register.mjs
similarity index 79%
rename from src/runtime/register.ts
rename to src/runtime/templates/register.mjs
index 0cd3854e..53d04826 100644
--- a/src/runtime/register.ts
+++ b/src/runtime/templates/register.mjs
@@ -2,3 +2,5 @@ import Vue from 'vue'
 import CompositionApi from '@vue/composition-api'
 
 Vue.use(CompositionApi)
+
+export default () => {}
diff --git a/test/fixture/nuxt.config.ts b/test/fixture/nuxt.config.ts
index 9a566198..899ef31f 100644
--- a/test/fixture/nuxt.config.ts
+++ b/test/fixture/nuxt.config.ts
@@ -10,25 +10,21 @@ const isTesting = process.env.NODE_ENV !== 'development'
 const rootDir = resolve(__dirname, '../..')
 
 const inDevelopment = !process.env.TEST_BUILT_MODULE
+const moduleSource = join(rootDir, inDevelopment ? 'src' : 'dist')
 
 console.log('Testing', inDevelopment ? 'source' : 'built', 'module')
 
 export default <NuxtConfig>{
   alias: {
-    '@nuxtjs/composition-api/dist/runtime/register': join(
-      rootDir,
-      inDevelopment ? 'src' : 'dist',
-      'runtime/register'
-    ),
     '@nuxtjs/composition-api/dist/runtime/globals': join(
-      rootDir,
-      inDevelopment ? 'src' : 'dist',
+      moduleSource,
       'runtime/globals'
     ),
-    '@nuxtjs/composition-api': join(
-      rootDir,
-      inDevelopment ? 'src/runtime/index.ts' : 'dist/runtime/index.js'
+    '@nuxtjs/composition-api/dist/babel-plugin': join(
+      moduleSource,
+      'babel-plugin'
     ),
+    '@nuxtjs/composition-api': join(moduleSource, 'runtime'),
   },
   target: isGenerated ? 'static' : 'server',
   publicRuntimeConfig: {
@@ -76,7 +72,7 @@ export default <NuxtConfig>{
   buildModules: [
     '@nuxt/typescript-build',
     '@nuxtjs/pwa',
-    join(rootDir, inDevelopment ? 'src' : 'dist', 'module'),
+    join(moduleSource, 'module'),
   ],
   pwa: {
     icon: false,
diff --git a/tsconfig.json b/tsconfig.json
index 4ef83cb2..f38b53aa 100644
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -15,9 +15,6 @@
       "@nuxtjs/composition-api/dist/runtime/globals": [
         "./src/runtime/globals.ts"
       ],
-      "@nuxtjs/composition-api/dist/runtime/register": [
-        "./src/runtime/register.ts"
-      ],
       "@nuxtjs/composition-api": ["./src/index.ts"]
     },
     "types": ["node", "@nuxt/types", "jest"]