refactor(appkit): own asUser OBO wrapping in Plugin proxy (#385)

Author: MarioCadenas

packages/appkit/src/core/appkit.ts

@@ -10,14 +10,13 @@ import type {
 } from "shared";
 import { version as productVersion } from "../../package.json";
 import { CacheManager } from "../cache";
-import { runInUserContext, ServiceContext } from "../context";
-import type { UserContext } from "../context/user-context";
+import { ServiceContext } from "../context";
 import {
   isInternalTelemetryEnabled,
   TelemetryReporter,
 } from "../internal-telemetry";
 import { createLogger } from "../logging/logger";
-import { USER_CONTEXT_SYMBOL } from "../plugin/plugin";
+import { isPlainObject } from "../plugin/plugin";
 import { ResourceRegistry, ResourceType } from "../registry";
 import type { TelemetryConfig } from "../telemetry";
 import { TelemetryManager } from "../telemetry";
@@ -128,45 +127,24 @@ export class AppKit<TPlugins extends InputPluginMap> {
       const val = exports[key];
       if (typeof val === "function") {
         exports[key] = (val as (...args: unknown[]) => unknown).bind(context);
-      } else if (AppKit.isPlainObject(val)) {
+      } else if (isPlainObject(val)) {
         this.bindExportMethods(val as Record<string, unknown>, context);
       }
     }
   }
 
-  /**
-   * Wraps all function properties in an exports object so they run
-   * inside the given user context (via AsyncLocalStorage).
-   * This ensures RoutingPool and other context-aware code sees the
-   * user identity even though the function was obtained outside the proxy.
-   */
-  private wrapExportsInUserContext(
-    exports: Record<string, unknown>,
-    userContext: UserContext,
-  ) {
-    for (const key in exports) {
-      if (!Object.hasOwn(exports, key)) continue;
-      const val = exports[key];
-      if (typeof val === "function") {
-        const fn = val as (...args: unknown[]) => unknown;
-        exports[key] = (...args: unknown[]) =>
-          runInUserContext(userContext, () => fn(...args));
-      } else if (AppKit.isPlainObject(val)) {
-        this.wrapExportsInUserContext(
-          val as Record<string, unknown>,
-          userContext,
-        );
-      }
-    }
-  }
-
   /**
    * Wraps a plugin's exports with an `asUser` method that returns
    * a user-scoped version of the exports.
    *
    * When `exports()` returns a callable (function), it is returned as-is
    * since the plugin manages its own `asUser` per-call (e.g. files plugin).
    * When it returns a plain object, the standard `asUser` wrapper is added.
+   *
+   * The OBO-side wrapping lives inside `Plugin.asUser` — calling
+   * `plugin.asUser(req).exports()` returns exports whose functions already
+   * run inside the user's AsyncLocalStorage scope. AppKit only adapts the
+   * shape; it does not own the user-context concept.
    */
   private wrapWithAsUser<T extends BasePlugin>(plugin: T) {
     // If plugin doesn't implement exports(), return empty object
@@ -192,40 +170,11 @@ export class AppKit<TPlugins extends InputPluginMap> {
        * Returns user-scoped exports where all methods execute with the
        * user's Databricks credentials instead of the service principal.
        */
-      asUser: (req: import("express").Request) => {
-        const userPlugin = (plugin as any).asUser(req);
-        const userContext = (userPlugin as any)[
-          USER_CONTEXT_SYMBOL
-        ] as UserContext;
-        const userExports = (plugin.exports?.() ?? {}) as Record<
-          string,
-          unknown
-        >;
-        // Wrap each export in runInUserContext instead of bind.
-        // bind() bypasses the Proxy get trap, so methods called via bind
-        // would not run inside the user's AsyncLocalStorage context.
-        if (userContext) {
-          this.wrapExportsInUserContext(userExports, userContext);
-        } else {
-          // Fallback for dev mode proxy (no userContext symbol)
-          this.bindExportMethods(userExports, userPlugin);
-        }
-        return userExports;
-      },
+      asUser: (req: import("express").Request) =>
+        (plugin as any).asUser(req).exports() as Record<string, unknown>,
     };
   }
 
-  /**
-   * Returns true if the value is a plain object (not an array, Date, etc.).
-   */
-  private static isPlainObject(
-    value: unknown,
-  ): value is Record<string, unknown> {
-    if (typeof value !== "object" || value === null) return false;
-    const proto = Object.getPrototypeOf(value);
-    return proto === Object.prototype || proto === null;
-  }
-
   static async _createApp<
     T extends PluginData<PluginConstructor, unknown, string>[],
   >(

packages/appkit/src/plugin/plugin.ts

@@ -14,12 +14,7 @@ import type {
 } from "shared";
 import { AppManager } from "../app";
 import { CacheManager } from "../cache";
-import {
-  getCurrentUserId,
-  runInUserContext,
-  ServiceContext,
-  type UserContext,
-} from "../context";
+import { getCurrentUserId, runInUserContext, ServiceContext } from "../context";
 import type { PluginContext } from "../core/plugin-context";
 import { AppKitError, AuthenticationError } from "../errors";
 import { createLogger } from "../logging/logger";
@@ -43,19 +38,55 @@ import type {
 
 const logger = createLogger("plugin");
 
-/**
- * Symbol used to expose the UserContext from an asUser() proxy.
- * Allows wrapWithAsUser in appkit.ts to retrieve the context and
- * wrap export methods in runInUserContext().
- */
-export const USER_CONTEXT_SYMBOL = Symbol("appkit.userContext");
-
 /**
  * OTel context key for marking OBO dev mode fallback.
  * Set when asUser() is called in development mode without a user token.
  */
 const DEV_OBO_FALLBACK_KEY = createContextKey("appkit.devOboFallback");
 
+/**
+ * Returns true if `value` is a plain object literal (not an array, Date,
+ * class instance, etc.). Used to decide whether to recurse into nested
+ * export shapes when wrapping functions.
+ *
+ * @internal exported so the AppKit core can reuse the same predicate for
+ * its `bindExportMethods` walk; not part of the public package surface.
+ */
+export function isPlainObject(
+  value: unknown,
+): value is Record<string, unknown> {
+  if (typeof value !== "object" || value === null) return false;
+  const proto = Object.getPrototypeOf(value);
+  return proto === Object.prototype || proto === null;
+}
+
+/**
+ * Returns a deep copy of `exports` where every function has been replaced
+ * with `wrap(fn)`, walking into nested plain objects.
+ *
+ * Used by the asUser proxy to make the user context follow function
+ * references that escape the proxy via `exports()`. The original input is
+ * not mutated, so plugins that memoize `exports()` are safe — each call
+ * through the proxy yields an independent, freshly wrapped view.
+ */
+function wrapExportFunctions(
+  exports: Record<string, unknown>,
+  wrap: (fn: (...a: unknown[]) => unknown) => (...a: unknown[]) => unknown,
+): Record<string, unknown> {
+  const result: Record<string, unknown> = {};
+  for (const key of Object.keys(exports)) {
+    const val = exports[key];
+    if (typeof val === "function") {
+      result[key] = wrap(val as (...a: unknown[]) => unknown);
+    } else if (isPlainObject(val)) {
+      result[key] = wrapExportFunctions(val, wrap);
+    } else {
+      result[key] = val;
+    }
+  }
+  return result;
+}
+
 /**
  * Returns true if the current execution is an OBO dev mode fallback
  * (asUser() was called but fell back to service principal due to missing token).
@@ -403,30 +434,18 @@ export abstract class Plugin<
     const userEmail = req.header("x-forwarded-email");
     const isDev = process.env.NODE_ENV === "development";
 
-    // In local development, skip user impersonation
-    // since there's no user token available
+    // In local development, skip user impersonation since there's no user
+    // token available. Mark execution as OBO dev fallback via OTel context
+    // so telemetry can distinguish intended OBO calls from regular SP calls.
     if (!token && isDev) {
       logger.warn(
         "asUser() called without user token in development mode. Skipping user impersonation.",
       );
 
-      // Return a proxy that marks execution as OBO dev fallback via OTel context,
-      // so telemetry spans can distinguish intended OBO calls from regular SP calls
-      return new Proxy(this, {
-        get: (target, prop, receiver) => {
-          const value = Reflect.get(target, prop, receiver);
-          if (typeof value !== "function") return value;
-          if (typeof prop === "string" && EXCLUDED_FROM_PROXY.has(prop))
-            return value;
-
-          return (...args: unknown[]) => {
-            const ctx = otelContext
-              .active()
-              .setValue(DEV_OBO_FALLBACK_KEY, true);
-            return otelContext.with(ctx, () => value.apply(target, args));
-          };
-        },
-      }) as this;
+      return this._createAsUserProxy((fn) => (...args) => {
+        const ctx = otelContext.active().setValue(DEV_OBO_FALLBACK_KEY, true);
+        return otelContext.with(ctx, () => fn(...args));
+      });
     }
 
     if (!token) {
@@ -446,34 +465,55 @@ export abstract class Plugin<
       userEmail ?? undefined,
     );
 
-    // Return a proxy that wraps method calls in user context
-    return this._createUserContextProxy(userContext);
+    return this._createAsUserProxy(
+      (fn) =>
+        (...args) =>
+          runInUserContext(userContext, () => fn(...args)),
+    );
   }
 
   /**
-   * Creates a proxy that wraps method calls in a user context.
-   * This allows all plugin methods to automatically use the user's
-   * Databricks credentials.
+   * Creates a proxy of `this` where every method call — and every function
+   * in the result of `exports()` — runs inside `wrapCall`.
+   *
+   * `wrapCall` decides the per-call scope. Two strategies are used today:
+   *   - real OBO:     fn => (...args) => runInUserContext(userContext, () => fn(...args))
+   *   - dev fallback: fn => (...args) => otelContext.with(DEV_OBO_FALLBACK_KEY=true, () => fn(...args))
+   *
+   * `exports` is intercepted because methods captured in the returned
+   * exports object never re-enter the proxy's `get` trap. Wrapping them
+   * here is the only way to make the user context follow function
+   * references back out of the plugin.
    */
-  private _createUserContextProxy(userContext: UserContext): this {
+  private _createAsUserProxy(
+    wrapCall: (
+      fn: (...a: unknown[]) => unknown,
+    ) => (...a: unknown[]) => unknown,
+  ): this {
     return new Proxy(this, {
       get: (target, prop, receiver) => {
-        // Expose userContext via symbol so wrapWithAsUser can wrap exports
-        if (prop === USER_CONTEXT_SYMBOL) return userContext;
-
         const value = Reflect.get(target, prop, receiver);
 
-        if (typeof value !== "function") {
+        if (typeof value !== "function") return value;
+        if (typeof prop === "string" && EXCLUDED_FROM_PROXY.has(prop))
           return value;
-        }
 
-        if (typeof prop === "string" && EXCLUDED_FROM_PROXY.has(prop)) {
-          return value;
+        if (prop === "exports") {
+          return () => {
+            const raw = (value as () => unknown).call(target);
+            if (raw == null) return {};
+            // Callable exports (e.g. files, jobs) manage per-call asUser
+            // themselves; leave them untouched.
+            if (typeof raw === "function") return raw;
+            if (isPlainObject(raw)) {
+              return wrapExportFunctions(raw, wrapCall);
+            }
+            return raw;
+          };
         }
 
-        return (...args: unknown[]) => {
-          return runInUserContext(userContext, () => value.apply(target, args));
-        };
+        const fn = (value as (...a: unknown[]) => unknown).bind(target);
+        return wrapCall(fn);
       },
     }) as this;
   }