From 869f3fb93e61400be4fd925e0850c2b1564749e2 Mon Sep 17 00:00:00 2001
From: Eduardo San Martin Morote <posva@users.noreply.github.com>
Date: Wed, 5 Apr 2023 09:18:13 +0200
Subject: [PATCH] feat(app): app.runWithContext() (#7451)

---
 .../__tests__/apiCreateApp.spec.ts            | 16 +++++++++++++
 packages/runtime-core/src/apiCreateApp.ts     | 23 +++++++++++++++++++
 packages/runtime-core/src/apiInject.ts        | 12 ++++++----
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/packages/runtime-core/__tests__/apiCreateApp.spec.ts b/packages/runtime-core/__tests__/apiCreateApp.spec.ts
index 7cfce51c161..699d0be6d50 100644
--- a/packages/runtime-core/__tests__/apiCreateApp.spec.ts
+++ b/packages/runtime-core/__tests__/apiCreateApp.spec.ts
@@ -110,6 +110,22 @@ describe('api: createApp', () => {
     expect(`App already provides property with key "bar".`).toHaveBeenWarned()
   })
 
+  test('runWithContext', () => {
+    const app = createApp({
+      setup() {
+        provide('foo', 'should not be seen')
+        return () => h('div')
+      }
+    })
+    app.provide('foo', 1)
+
+    expect(app.runWithContext(() => inject('foo'))).toBe(1)
+
+    // ensure the context is restored
+    inject('foo')
+    expect('inject() can only be used inside setup').toHaveBeenWarned()
+  })
+
   test('component', () => {
     const Root = {
       // local override
diff --git a/packages/runtime-core/src/apiCreateApp.ts b/packages/runtime-core/src/apiCreateApp.ts
index 05c7ce31539..6b698ba26c2 100644
--- a/packages/runtime-core/src/apiCreateApp.ts
+++ b/packages/runtime-core/src/apiCreateApp.ts
@@ -51,6 +51,14 @@ export interface App<HostElement = any> {
   unmount(): void
   provide<T>(key: InjectionKey<T> | string, value: T): this
 
+  /**
+   * Runs a function with the app as active instance. This allows using of `inject()` within the function to get access
+   * to variables provided via `app.provide()`.
+   *
+   * @param fn - function to run with the app as active instance
+   */
+  runWithContext<T>(fn: () => T): T
+
   // internal, but we need to expose these for the server-renderer and devtools
   _uid: number
   _component: ConcreteComponent
@@ -370,6 +378,15 @@ export function createAppAPI<HostElement>(
         context.provides[key as string | symbol] = value
 
         return app
+      },
+
+      runWithContext(fn) {
+        currentApp = app
+        try {
+          return fn()
+        } finally {
+          currentApp = null
+        }
       }
     })
 
@@ -380,3 +397,9 @@ export function createAppAPI<HostElement>(
     return app
   }
 }
+
+/**
+ * @internal Used to identify the current app when using `inject()` within
+ * `app.runWithContext()`.
+ */
+export let currentApp: App<unknown> | null = null
diff --git a/packages/runtime-core/src/apiInject.ts b/packages/runtime-core/src/apiInject.ts
index c5c47876cb3..6eedee88c09 100644
--- a/packages/runtime-core/src/apiInject.ts
+++ b/packages/runtime-core/src/apiInject.ts
@@ -1,6 +1,7 @@
 import { isFunction } from '@vue/shared'
 import { currentInstance } from './component'
 import { currentRenderingInstance } from './componentRenderContext'
+import { currentApp } from './apiCreateApp'
 import { warn } from './warning'
 
 export interface InjectionKey<T> extends Symbol {}
@@ -46,21 +47,24 @@ export function inject(
   // fallback to `currentRenderingInstance` so that this can be called in
   // a functional component
   const instance = currentInstance || currentRenderingInstance
-  if (instance) {
+
+  // also support looking up from app-level provides w/ `app.runWithContext()`
+  if (instance || currentApp) {
     // #2400
     // to support `app.use` plugins,
     // fallback to appContext's `provides` if the instance is at root
-    const provides =
-      instance.parent == null
+    const provides = instance
+      ? instance.parent == null
         ? instance.vnode.appContext && instance.vnode.appContext.provides
         : instance.parent.provides
+      : currentApp!._context.provides
 
     if (provides && (key as string | symbol) in provides) {
       // TS doesn't allow symbol as index type
       return provides[key as string]
     } else if (arguments.length > 1) {
       return treatDefaultAsFactory && isFunction(defaultValue)
-        ? defaultValue.call(instance.proxy)
+        ? defaultValue.call(instance && instance.proxy)
         : defaultValue
     } else if (__DEV__) {
       warn(`injection "${String(key)}" not found.`)