diff --git a/docs/docs/plugins/execution-context.md b/docs/docs/plugins/execution-context.md index 1764bf7c5..9ddbd9f05 100644 --- a/docs/docs/plugins/execution-context.md +++ b/docs/docs/plugins/execution-context.md @@ -114,6 +114,31 @@ is already open, fallback retains it instead of widening to SP. The marker does not leak outside the scope. Production never falls back when credentials are missing. +## App-only resources and missing credentials + +The manifest capability contract keeps `secret`, `database`, and `postgres` +app-only for the new `appkit.asUser` and `runInCallerContext` APIs. Accessing +those resources through these caller-scoped plugin APIs or tools +raises a clear error identifying the resource by its manifest alias, or its type +when no alias is available. The message states that the resource is app-only in +this version of AppKit and does not support OBO execution through these APIs. +The check applies when using cached handles inside a later user scope too. +It does not reject unrelated plugins merely because an app-only plugin is +installed. Required resources, runtime requirements, and configured optional +resources determine the plugin's resource capability. + +Deprecated `plugin.asUser` and `runInUserContext`, direct request-based tool +dispatch, and the existing agents HTTP routes retain their established resource +behavior, including Lakebase per-user routing. They still establish user identity +and reject missing production credentials; they never fall back to SP by omission. +Using a deprecated entry point inside a new guarded scope cannot disable its guards. + +Missing-token messages distinguish OBO-capable resources from generic operations. +For an OBO-capable resource, the message explains that no user token was forwarded +and the app may be deployed service-principal-only. Otherwise the generic missing +user token error remains. The app-level check uses the registered plugins' active +resource metadata because no individual plugin has been selected yet. + ## Credential expiration and telemetry A structured downstream HTTP 401 inside a caller scope throws diff --git a/docs/docs/plugins/lakebase.md b/docs/docs/plugins/lakebase.md index d3f77e78d..cf1fa9cb2 100644 --- a/docs/docs/plugins/lakebase.md +++ b/docs/docs/plugins/lakebase.md @@ -113,88 +113,28 @@ await createApp({ }); ``` -## On-Behalf-Of (OBO) — per-user connections +## On-Behalf-Of (OBO) {#on-behalf-of-obo--per-user-connections} -When your app needs Row-Level Security (RLS) or per-user data isolation, use `asUser(req)` to execute queries using a per-user Lakebase connection pool. Each user's pool is authenticated with their Databricks identity, so PostgreSQL's `current_user` reflects the actual user. +Lakebase is app-only through the new `appkit.asUser` and `runInCallerContext` +APIs. Operations in those scopes fail with a clear error. For a connector call +without a manifest alias: -### Prerequisites - -1. **Enable user authorization** in your Databricks App with the **`postgres`** scope. See [User authorization](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth#user-authorization) for setup instructions. In your `databricks.yml`: - ```yaml - resources: - apps: - app: - user_api_scopes: - - postgres - ``` - Apps scaffolded with `databricks apps init` and the Lakebase plugin include this automatically. - -2. Each app user needs a **Postgres role** in Lakebase. Create one with the Databricks CLI: - - ```bash - databricks postgres create-role "projects/{project_id}/branches/{branch_id}" \ - --json '{"spec": {"identity_type": "USER", "postgres_role": "user@example.com"}}' - ``` - - Alternatively, create roles in the Lakebase UI under **Branch Overview** → **Add role**. - - :::note - Do not grant `databricks_superuser` to OBO users — superusers bypass RLS. Use [fine-grained grants](#fine-grained-permissions) instead. - ::: - -### Usage - -No configuration needed — just call `asUser(req)`: - -```ts -const AppKit = await createApp({ - plugins: [server(), lakebase()], -}); - -// Service principal query (default — bypasses RLS as table owner) -const all = await AppKit.lakebase.query("SELECT * FROM app.orders"); - -// User-scoped query (per-user pool, RLS enforced) -app.get("/api/my-orders", async (req, res) => { - const result = await AppKit.lakebase - .asUser(req) - .query("SELECT * FROM app.orders ORDER BY created_at DESC"); - res.json(result.rows); -}); +```text +Resource "postgres" is app-only in this version of AppKit and does not support OBO (on-behalf-of-user) execution. It runs as the service principal. Resource type: postgres. ``` -When `asUser(req)` is called: -1. The user's token and identity are extracted from `x-forwarded-access-token` and `x-forwarded-email` headers (set automatically by Databricks Apps). -2. A per-user `pg.Pool` is created (or reused) with the user's OAuth credentials. -3. `query()` and `pool` use the user's pool — `current_user` in PostgreSQL reflects the user's identity. - -### Row-Level Security example - -```sql --- As the service principal (during app setup): -ALTER TABLE app.orders ENABLE ROW LEVEL SECURITY; - -CREATE POLICY user_orders ON app.orders - FOR ALL TO PUBLIC - USING (owner = current_user); - --- Grant access so OBO users can query -GRANT USAGE ON SCHEMA app TO PUBLIC; -GRANT SELECT, INSERT ON ALL TABLES IN SCHEMA app TO PUBLIC; -``` - -### How it works - -- The **service principal pool** (`AppKit.lakebase.pool`) is always created and used for DDL operations, seeding, and admin queries. -- **Per-user pools** are created on the first `asUser(req)` call and cached by user identity. Each pool has its own OAuth token refresh cycle. -- Idle connections within per-user pools close automatically (30s idle timeout). Empty pool objects are cleaned up periodically. -- On shutdown, all pools (SP + user) are closed gracefully. -- In development mode (`NODE_ENV=development`), if no user token is available, `asUser(req)` falls back to the SP pool with a warning. - -:::caution[RLS and superusers] -PostgreSQL superusers bypass Row-Level Security entirely. Users with the `databricks_superuser` role will see all rows regardless of RLS policies. For RLS enforcement, use [fine-grained grants](#fine-grained-permissions) instead of the superuser role. -::: - +Use a plain Lakebase call outside a caller scope for SP execution. There is no +`asApp()` escape from a caller scope. The guard also applies to agent tools, ORM +configuration, and pool handles obtained before entering a new caller scope. + +The platform has a `postgres` user API scope, but AppKit does not enable that +capability in the new v1 execution API. For backward compatibility, deprecated +`plugin.asUser(req)` and `runInUserContext` retain the existing per-user pool +routing, as do direct request-based tool dispatch and existing agents routes. +These paths keep the user's identity and do not silently use the SP. They cannot +disable an enclosing new caller scope's guard. Existing applications can upgrade +without migrating their Lakebase OBO calls, then adopt the new execution API +explicitly. Database permissions and row-level policies still govern data access. ## Database Permissions When you create the app with the Lakebase resource using the [Getting started](#getting-started-with-the-lakebase) guide, the Service Principal is automatically granted `CONNECT_AND_CREATE` permission on the `postgres` resource. This lets the Service Principal connect to the database and create new objects, but **not access any existing schemas or tables.** diff --git a/packages/appkit/src/connectors/lakebase/routing-pool.ts b/packages/appkit/src/connectors/lakebase/routing-pool.ts index 1a8ad41d6..e092403d6 100644 --- a/packages/appkit/src/connectors/lakebase/routing-pool.ts +++ b/packages/appkit/src/connectors/lakebase/routing-pool.ts @@ -2,6 +2,7 @@ import type { Pool, PoolClient, QueryResult, QueryResultRow } from "pg"; import type { CallerContext } from "../../context/caller-context"; import { getCallerContext } from "../../context/execution-context"; +import { assertResourceExecution } from "../../context/resource-capabilities"; /** * Subset of `pg.Pool` exposed by the Lakebase plugin. @@ -23,16 +24,8 @@ export interface LakebasePool { } /** - * A `pg.Pool`-like wrapper that routes queries to the appropriate pool - * based on the current execution context. - * - * When called inside `runInCallerContext()` (set up by `Plugin.asUser(req)`), - * queries route to the per-user pool returned by `resolveUserPool`. - * Otherwise, queries route to the service-principal pool. - * - * This enables OBO (On-Behalf-Of) without custom `asUser()` overrides — - * the base class sets up AsyncLocalStorage context, and the RoutingPool - * reads it transparently. + * A `pg.Pool`-like compatibility wrapper. New caller scopes enforce the v1 + * app-only contract. Deprecated entry points retain existing per-user routing. */ export class RoutingPool implements LakebasePool { constructor( @@ -41,6 +34,8 @@ export class RoutingPool implements LakebasePool { ) {} private activePool(): Pool { + // Strict scopes reject before resolving a pool; legacy OBO keeps its user pool. + assertResourceExecution("postgres"); const userCtx = getCallerContext(); return userCtx ? this.resolveUserPool(userCtx) : this.spPool; } diff --git a/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts b/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts index 29912b098..80a824b52 100644 --- a/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts @@ -1,6 +1,7 @@ import type { Pool } from "pg"; import { describe, expect, test, vi } from "vitest"; +import { runInCallerContext } from "../../../context/execution-context"; import { RoutingPool } from "../routing-pool"; function makeMockPool(label: string) { @@ -18,6 +19,28 @@ function makeMockPool(label: string) { } describe("RoutingPool", () => { + test("new caller scopes reject queries and connects without touching either pool", () => { + const spPool = makeMockPool("sp"); + const userPool = makeMockPool("user"); + const resolveUserPool = vi.fn(() => userPool); + const pool = new RoutingPool(spPool, resolveUserPool); + const caller = { + client: {} as any, + principal: { type: "user" as const, userId: "alice" }, + workspaceId: Promise.resolve("workspace"), + }; + for (const operation of [ + () => pool.query("SELECT 1"), + () => pool.connect(), + ]) { + expect(() => + runInCallerContext>(caller, operation), + ).toThrow('Resource "postgres" is app-only in this version of AppKit'); + } + expect(resolveUserPool).not.toHaveBeenCalled(); + expect(spPool.query).not.toHaveBeenCalled(); + expect(spPool.connect).not.toHaveBeenCalled(); + }); test("routes to SP pool when no user context is active", async () => { const spPool = makeMockPool("sp"); const userPool = makeMockPool("user"); diff --git a/packages/appkit/src/context/execution-context.ts b/packages/appkit/src/context/execution-context.ts index 9f87ddbf4..e3ab8e3c2 100644 --- a/packages/appkit/src/context/execution-context.ts +++ b/packages/appkit/src/context/execution-context.ts @@ -19,16 +19,28 @@ import { type UserContext, } from "./user-context"; -const executionContextStorage = new AsyncLocalStorage(); +interface CallerScope { + readonly caller: CallerContext; + readonly legacyExecution: boolean; +} + +const executionContextStorage = new AsyncLocalStorage(); function runInCallerScope( callerContext: CallerContext, fn: () => T, legacyResources?: WarehouseBinding, + legacyExecution = false, ): T { - const caller = snapshotCallerContext(callerContext); + const scope = Object.freeze({ + caller: snapshotCallerContext(callerContext), + // Deprecated entry points cannot disable guards in an existing new scope. + legacyExecution: + legacyExecution && + (executionContextStorage.getStore()?.legacyExecution ?? true), + }); return runWithResourceBindings(legacyResources, () => - executionContextStorage.run(caller, () => { + executionContextStorage.run(scope, () => { try { const result = fn(); if (result instanceof Promise) { @@ -71,6 +83,19 @@ export function runInCallerContext( return runInCallerScope(callerContext, fn); } +/** @internal Preserve legacy resource behavior without weakening an enclosing scope. */ +export function runInLegacyCallerContext( + caller: CallerContext, + fn: () => T, +): T { + return runInCallerScope(caller, fn, undefined, true); +} + +/** @internal Legacy user entry points retain their established resource routing. */ +export function isLegacyCallerScope(): boolean { + return executionContextStorage.getStore()?.legacyExecution === true; +} + /** @deprecated Use runInCallerContext. */ export function runInUserContext( userContext: UserContext | (CallerContext & Pick), @@ -82,9 +107,10 @@ export function runInUserContext( toCallerContext(userContext), fn, Object.freeze({ warehouseId: userContext.warehouseId }), + true, ); } - return runInCallerContext(toCallerContext(userContext), fn); + return runInLegacyCallerContext(toCallerContext(userContext), fn); } /** @@ -98,7 +124,7 @@ export function runInUserContext( export function getExecutionContext(): | ServiceContextState | (CallerContext & UserContext) { - const callerContext = executionContextStorage.getStore(); + const callerContext = getCallerContext(); if (callerContext) { return legacyUserContext(callerContext, captureWarehouseId()); } @@ -118,6 +144,11 @@ export function getCurrentActorId(): string | undefined { return getCallerContext()?.principal.userId; } +/** @internal Bare ID of the effective caller or service principal. */ +export function getCurrentPrincipalId(): string { + return getCurrentActorId() ?? ServiceContext.get().serviceUserId; +} + /** * @deprecated Use getCurrentPrincipalKey for new cache keys or getCurrentActorId * for audit. Preserves the bare user or service ID for existing callers. @@ -127,7 +158,7 @@ export function getCurrentUserId(): string { "getCurrentUserId", "getCurrentPrincipalKey (cache) or getCurrentActorId (audit)", ); - return getCurrentActorId() ?? ServiceContext.get().serviceUserId; + return getCurrentPrincipalId(); } /** @@ -169,12 +200,12 @@ export function isInUserContext(): boolean { * to be initialized and never throws. */ export function getCallerContext(): CallerContext | undefined { - return executionContextStorage.getStore(); + return executionContextStorage.getStore()?.caller; } /** @deprecated Use getCallerContext and its principal field. */ export function getUserContext(): (CallerContext & UserContext) | undefined { warnContextDeprecation("getUserContext", "getCallerContext"); - const scope = executionContextStorage.getStore(); - return scope ? legacyUserContext(scope, captureWarehouseId()) : undefined; + const caller = getCallerContext(); + return caller ? legacyUserContext(caller, captureWarehouseId()) : undefined; } diff --git a/packages/appkit/src/context/index.ts b/packages/appkit/src/context/index.ts index c9f213bb8..b3958d0af 100644 --- a/packages/appkit/src/context/index.ts +++ b/packages/appkit/src/context/index.ts @@ -1,6 +1,7 @@ export { getCallerContext, getCurrentActorId, + getCurrentPrincipalId, getCurrentPrincipalKey, getCurrentUserId, getExecutionContext, diff --git a/packages/appkit/src/context/request-scope.ts b/packages/appkit/src/context/request-scope.ts index 8c58ba17c..f95210397 100644 --- a/packages/appkit/src/context/request-scope.ts +++ b/packages/appkit/src/context/request-scope.ts @@ -3,7 +3,11 @@ import type { Request } from "express"; import { AuthenticationError } from "../errors"; import { createLogger } from "../logging/logger"; -import { runInCallerContext } from "./execution-context"; +import { + runInCallerContext, + runInLegacyCallerContext, +} from "./execution-context"; +import { hasOboResource } from "./resource-capabilities"; import { ServiceContext } from "./service-context"; const logger = createLogger("execution-context"); @@ -21,7 +25,11 @@ export function isDevOboFallback(): boolean { /** Build the caller once from the trusted Apps proxy headers. */ export function createRequestScope( req: Request, - createCaller = ServiceContext.createCallerContext, + resourceTypes: readonly string[] = [], + options: { + createCaller?: typeof ServiceContext.createCallerContext; + legacy?: boolean; + } = {}, ): RequestScope { const token = req.header("x-forwarded-access-token")?.trim(); const userId = req.header("x-forwarded-user")?.trim(); @@ -40,14 +48,26 @@ export function createRequestScope( ), }; } - if (!token) throw AuthenticationError.missingToken("user token"); + if (!token) { + if (hasOboResource(resourceTypes)) { + const message = + "This resource is OBO-capable but no user token was forwarded. The app is likely deployed service-principal-only. Enable user authorization and forward x-forwarded-access-token."; + throw new AuthenticationError(message, { clientMessage: message }); + } + throw AuthenticationError.missingToken("user token"); + } if (!userId && !isDev) throw AuthenticationError.missingUserId(); - const caller = createCaller( + const caller = (options.createCaller ?? ServiceContext.createCallerContext)( token, userId || "dev-user", undefined, userEmail, ); - return { run: (fn) => runInCallerContext(caller, fn) }; + return { + run: (fn) => + options.legacy + ? runInLegacyCallerContext(caller, fn) + : runInCallerContext(caller, fn), + }; } diff --git a/packages/appkit/src/context/resource-capabilities.ts b/packages/appkit/src/context/resource-capabilities.ts new file mode 100644 index 000000000..8b807c952 --- /dev/null +++ b/packages/appkit/src/context/resource-capabilities.ts @@ -0,0 +1,93 @@ +import { + APP_ONLY_RESOURCE_TYPES, + SCOPE_BY_TYPE, + type BasePluginConfig, + type PluginConstructor, +} from "shared"; + +import { AppKitError } from "../errors/base"; +import { getCallerContext, isLegacyCallerScope } from "./execution-context"; +import { scopeApi } from "./scoped-api"; + +class AppOnlyResourceError extends AppKitError { + readonly code = "APP_ONLY_RESOURCE"; + readonly statusCode = 400; + readonly isRetryable = false; + constructor(type: string, name = type) { + const message = `Resource "${name}" is app-only in this version of AppKit and does not support OBO (on-behalf-of-user) execution. It runs as the service principal. Resource type: ${type}.`; + super(message, { clientMessage: message, context: { resourceType: type } }); + } +} + +/** Read the same required/bound resource metadata used by provisioning. */ +function getPluginResources(plugin: object) { + const ctor = plugin.constructor as Partial; + const config = (plugin as { config?: BasePluginConfig }).config; + const resources = ctor.manifest?.resources; + const optional = + resources?.optional?.filter((resource) => + Object.values(resource.fields).some( + (field) => field.env && process.env[field.env]?.trim(), + ), + ) ?? []; + const runtime = + config && ctor.getResourceRequirements + ? ctor.getResourceRequirements(config) + : []; + return [ + ...(resources?.required ?? []), + ...optional, + ...runtime.filter((resource) => resource.required), + ]; +} + +export function getPluginResourceTypes(plugin: object): string[] { + return [ + ...new Set(getPluginResources(plugin).map((resource) => resource.type)), + ]; +} + +export function hasOboResource(types: readonly string[]): boolean { + return types.some((type) => Object.hasOwn(SCOPE_BY_TYPE, type)); +} + +function isAppOnly(type: string): boolean { + return (APP_ONLY_RESOURCE_TYPES as ReadonlySet).has(type); +} + +export function assertResourceExecution(type: string): void { + if (getCallerContext() && !isLegacyCallerScope() && isAppOnly(type)) + throw new AppOnlyResourceError(type); +} + +export function assertPluginExecution( + plugin: object, + callerRequested = false, +): void { + if (!callerRequested && (!getCallerContext() || isLegacyCallerScope())) + return; + for (const resource of getPluginResources(plugin)) { + if (isAppOnly(resource.type)) { + throw new AppOnlyResourceError( + resource.type, + resource.alias?.trim() || resource.type, + ); + } + } +} + +/** Cached app handles still check the principal at invocation time. */ +export function guardPluginApi(plugin: object, value: T): T { + if (!getPluginResourceTypes(plugin).some(isAppOnly)) return value; + return scopeApi( + value, + { + run: (fn) => { + assertPluginExecution(plugin); + return fn(); + }, + }, + plugin, + true, + ); +} diff --git a/packages/appkit/src/context/scoped-api.ts b/packages/appkit/src/context/scoped-api.ts index 143203617..7a6142ce7 100644 --- a/packages/appkit/src/context/scoped-api.ts +++ b/packages/appkit/src/context/scoped-api.ts @@ -56,22 +56,28 @@ export function scopeApi( value: T, scope: RequestScope, receiver?: unknown, + preserveIdentityMethods = false, ): T { if (typeof value === "function") { return new Proxy(value, { - apply: (fn, thisArg, args) => - scopeApi( - scope.run(() => Reflect.apply(fn, receiver ?? thisArg, args)), - scope, - ), + apply: (fn, thisArg, args) => { + const result = scope.run(() => + Reflect.apply(fn, receiver ?? thisArg, args), + ); + // Ambient resource guards must not change ordinary SP result objects. + if (preserveIdentityMethods) return result; + return scopeApi(result, scope); + }, get: (fn, key) => - isIdentityMethod(key) + !preserveIdentityMethods && isIdentityMethod(key) ? undefined - : scopeApi(Reflect.get(fn, key), scope, fn), + : scopeApi(Reflect.get(fn, key), scope, fn, preserveIdentityMethods), }); } if (value instanceof Promise) { - return value.then((result) => scopeApi(result, scope)) as T; + return value.then((result) => + scopeApi(result, scope, undefined, preserveIdentityMethods), + ) as T; } // Native streams pass through unchanged. The authenticated request already // ran inside the caller scope; reading the body is pure data with no @@ -113,7 +119,7 @@ export function scopeApi( if (!needsScope(value)) return value; const result = Object.create(Object.getPrototypeOf(value)); for (const key of Reflect.ownKeys(value)) { - if (isIdentityMethod(key)) continue; + if (!preserveIdentityMethods && isIdentityMethod(key)) continue; Object.defineProperty(result, key, { enumerable: Object.getOwnPropertyDescriptor(value, key)?.enumerable, configurable: true, @@ -121,7 +127,12 @@ export function scopeApi( const member = scope.run(() => Reflect.get(value, key)); // Data members are returned as-is so repeated reads are identical. return needsScope(member) - ? scopeApi(member, scope, receiver ?? value) + ? scopeApi( + member, + scope, + receiver ?? value, + preserveIdentityMethods, + ) : member; }, set: (next: unknown) => { diff --git a/packages/appkit/src/context/tests/resource-capabilities.test.ts b/packages/appkit/src/context/tests/resource-capabilities.test.ts new file mode 100644 index 000000000..e08a1bf0c --- /dev/null +++ b/packages/appkit/src/context/tests/resource-capabilities.test.ts @@ -0,0 +1,259 @@ +import type { AgentToolDefinition } from "shared"; +import { describe, expect, test, vi } from "vitest"; + +import { createAgent, runAgent } from "../../core/agent"; +import { Plugin, toPlugin } from "../../plugin"; +import { defineManifest } from "../../registry"; +import { + createMockRequest, + createTestApp, + createTestPluginContext, + createMockWorkspaceClient, +} from "../../testing"; +import { + getCurrentPrincipalKey, + runInCallerContext, + runInUserContext, +} from "../execution-context"; +import { + assertPluginExecution, + assertResourceExecution, +} from "../resource-capabilities"; +import { guardPluginApi } from "../resource-capabilities"; + +function probe(name: string, type?: string, alias = "resource") { + class ResourceProbe extends Plugin { + static manifest = defineManifest({ + name, + displayName: name, + description: "probe", + resources: { + required: type + ? [ + { + type, + alias, + resourceKey: "resource", + description: "probe", + permission: + type === "secret" + ? "READ" + : type === "sql_warehouse" + ? "CAN_USE" + : "CAN_CONNECT_AND_CREATE", + fields: { id: { description: "Test resource" } }, + }, + ] + : [], + optional: [], + }, + }); + getAgentTools(): AgentToolDefinition[] { + return [ + { name: "read", description: "probe", parameters: { type: "object" } }, + ]; + } + async executeAgentTool() { + return "read"; + } + exports() { + return { read: () => "read", nested: { read: () => "nested" } }; + } + } + return ResourceProbe; +} +const request = () => createMockRequest({ obo: { userId: "alice" } }); +const caller = { + principal: { type: "user" as const, userId: "alice" }, + client: createMockWorkspaceClient(), + workspaceId: Promise.resolve("workspace"), +}; + +describe("resource identity contract", () => { + test.each(["secret", "database", "postgres"])( + "labels %s errors from the manifest and preserves the error contract", + (type) => { + const AppOnly = probe("appOnly", type, "Orders storage"); + expect(() => assertPluginExecution(new AppOnly({}), true)).toThrow( + expect.objectContaining({ + message: `Resource "Orders storage" is app-only in this version of AppKit and does not support OBO (on-behalf-of-user) execution. It runs as the service principal. Resource type: ${type}.`, + code: "APP_ONLY_RESOURCE", + statusCode: 400, + isRetryable: false, + context: { resourceType: type }, + }), + ); + }, + ); + + test.each(["secret", "database", "postgres"])( + "uses the %s type as the label when no manifest is available", + (type) => { + expect(() => + runInCallerContext(caller, () => assertResourceExecution(type)), + ).toThrow(`Resource "${type}" is app-only in this version of AppKit`); + }, + ); + + test("resource guards preserve SP result identity", async () => { + const AppOnly = probe("appOnly", "postgres"); + const plugin = new AppOnly({}); + const rows = { rows: [{ id: 1 }] }; + const api = guardPluginApi(plugin, { + read: () => rows, + asyncRead: async () => rows, + }); + expect(api.read()).toBe(rows); + expect(await api.asyncRead()).toBe(rows); + }); + test.each(["secret", "database", "postgres"])( + "rejects %s in caller scope but preserves SP and unrelated plugins", + async (type) => { + const AppOnly = probe("appOnly", type); + const ObO = probe("analytics", "sql_warehouse"); + await using app = await createTestApp({ + plugins: [toPlugin(AppOnly)(), toPlugin(ObO)()], + server: false, + }); + const cachedRead = app.plugins.appOnly.read; + expect(cachedRead()).toBe("read"); + expect(app.plugins.appOnly.asUser(request()).read()).toBe("read"); + const scoped = app.plugins.asUser(request()); + expect(() => scoped.appOnly).toThrow(/does not support OBO/); + expect(scoped.analytics.read()).toBe("read"); + await expect(scoped.run(() => cachedRead())).rejects.toThrow( + /does not support OBO/, + ); + await expect( + scoped.run(() => app.plugins.appOnly.read()), + ).rejects.toThrow(/does not support OBO/); + await expect( + scoped.run(() => app.plugins.appOnly.asUser(request()).read()), + ).rejects.toThrow(/does not support OBO/); + await expect( + scoped.run(() => runInUserContext(caller, () => cachedRead())), + ).rejects.toThrow(/does not support OBO/); + expect(app.plugins.appOnly.read()).toBe("read"); + }, + ); + + test("keeps app-only, OBO-capable missing-token, and generic messages distinct", async () => { + const AppOnly = probe("appOnly", "postgres"); + const ObO = probe("analytics", "sql_warehouse"); + const Generic = probe("generic"); + await using app = await createTestApp({ + plugins: [toPlugin(AppOnly)(), toPlugin(ObO)(), toPlugin(Generic)()], + server: false, + }); + const missing = createMockRequest(); + expect(() => app.plugins.asUser(request()).appOnly).toThrow( + 'Resource "resource" is app-only in this version of AppKit', + ); + expect(() => app.plugins.appOnly.asUser(missing)).toThrow( + "Missing user token in request headers", + ); + expect(() => app.plugins.analytics.asUser(missing)).toThrow( + "OBO-capable but no user token was forwarded", + ); + expect(() => app.plugins.generic.asUser(missing)).toThrow( + "Missing user token in request headers", + ); + expect(() => app.plugins.asUser(missing)).toThrow( + "likely deployed service-principal-only", + ); + vi.stubEnv("NODE_ENV", "development"); + try { + expect(app.plugins.asUser(missing).analytics.read()).toBe("read"); + } finally { + vi.unstubAllEnvs(); + } + }); + + test("blocks direct and standalone toolkit dispatch before the provider is called", async () => { + const AppOnly = probe("appOnly", "postgres"); + const execute = vi.spyOn(AppOnly.prototype, "executeAgentTool"); + const mock = createTestPluginContext(); + await mock.attach(new AppOnly({})); + await expect( + runInCallerContext(caller, () => + mock.ctx.executeTool(request() as never, "appOnly", "read", {}), + ), + ).rejects.toThrow(/does not support OBO/); + const def = createAgent({ + instructions: "probe", + tools: (plugins) => plugins.appOnly.toolkit(), + model: { + async *run(_input, ctx) { + yield { + type: "message_delta", + content: String(await ctx.executeTool("appOnly.read", {})), + }; + }, + }, + }); + await expect( + runInCallerContext(caller, () => + runAgent(def, { messages: "hi", plugins: [toPlugin(AppOnly)()] }), + ), + ).rejects.toThrow(/does not support OBO/); + expect(execute).not.toHaveBeenCalled(); + execute.mockRestore(); + }); + + test("legacy direct dispatch preserves user identity for existing app-only providers", async () => { + const AppOnly = probe("appOnly", "postgres"); + const identity = vi + .spyOn(AppOnly.prototype, "executeAgentTool") + .mockImplementation(async () => getCurrentPrincipalKey()); + try { + const mock = createTestPluginContext(); + await mock.attach(new AppOnly({})); + await expect( + mock.ctx.executeTool(request(), "appOnly", "read", {}), + ).resolves.toBe("user:alice"); + identity.mockClear(); + await expect( + mock.ctx.executeTool(createMockRequest(), "appOnly", "read", {}), + ).rejects.toThrow(/token/i); + expect(identity).not.toHaveBeenCalled(); + } finally { + identity.mockRestore(); + } + }); + + test("does not treat an unbound optional app-only resource as active", () => { + class Optional extends Plugin { + static manifest = defineManifest({ + name: "optional", + displayName: "optional", + description: "probe", + resources: { + required: [], + optional: [ + { + type: "secret", + alias: "secret", + resourceKey: "secret", + description: "probe", + permission: "READ", + fields: { + key: { env: "APPKIT_TEST_OPTIONAL_SECRET", description: "key" }, + }, + }, + ], + }, + }); + } + const plugin = new Optional({}); + vi.stubEnv("APPKIT_TEST_OPTIONAL_SECRET", ""); + try { + expect(() => assertPluginExecution(plugin, true)).not.toThrow(); + vi.stubEnv("APPKIT_TEST_OPTIONAL_SECRET", "configured"); + expect(() => assertPluginExecution(plugin, true)).toThrow( + 'Resource "secret" is app-only in this version of AppKit', + ); + } finally { + vi.unstubAllEnvs(); + } + }); +}); diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index e0c4245ff..b1ce8efef 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -17,6 +17,7 @@ import { } from "../../agents/supervisor-api"; import { type CallerPrincipal, runInCallerContext } from "../../context"; import { getClientOptions } from "../../context/client-options"; +import { assertPluginExecution } from "../../context/resource-capabilities"; import { AuthenticationError, ConfigurationError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { createWorkspaceClient } from "../../workspace-client"; @@ -181,6 +182,7 @@ async function runAgentInternal( return entry.tool.execute(args as Record); } if (entry.kind === "toolkit") { + assertPluginExecution(entry.provider); return entry.provider.executeAgentTool( entry.localName, args as Record, diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 3a02fc2a0..7de60ce2a 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -12,6 +12,11 @@ import { version as productVersion } from "../../package.json"; import { CacheManager } from "../cache"; import { ServiceContext } from "../context"; import { createRequestScope } from "../context/request-scope"; +import { + assertPluginExecution, + getPluginResourceTypes, + guardPluginApi, +} from "../context/resource-capabilities"; import { scopeApi } from "../context/scoped-api"; import { isInternalTelemetryEnabled, @@ -47,7 +52,10 @@ export class AppKit { /** Execute a block or a single plugin call as the requesting user. */ asUser(req: import("express").Request) { - const scope = createRequestScope(req); + const scope = createRequestScope( + req, + Object.values(this.#pluginInstances).flatMap(getPluginResourceTypes), + ); const kit: Record = Object.create(null); // Memoize each plugin's scoped export for the life of THIS handle only. // The map is created per asUser(req) call and captured solely by this @@ -60,10 +68,16 @@ export class AppKit { enumerable: true, get: () => { if (!scopedExports.has(name)) { + // The app-only guard runs on first access; it is deterministic per + // (plugin, scope), so caching the result after it passes is safe, + // and if it throws nothing is cached (stays fail-closed). scopedExports.set( name, scopeApi( - scope.run(() => plugin.exports?.() ?? {}), + scope.run(() => { + assertPluginExecution(plugin, true); + return plugin.exports?.() ?? {}; + }), scope, plugin, ), @@ -204,12 +218,13 @@ export class AppKit { * to the shared request scope; new code should use the app-level API. */ private wrapWithAsUser(plugin: T) { + assertPluginExecution(plugin); // If plugin doesn't implement exports(), return empty object const pluginExports = plugin.exports?.() ?? {}; // If exports is a function, the plugin manages its own asUser pattern if (typeof pluginExports === "function") { - return pluginExports; + return guardPluginApi(plugin, pluginExports); } const objExports = pluginExports as Record; @@ -217,10 +232,10 @@ export class AppKit { // If plugin doesn't support asUser (no asUser method), return exports as-is if (typeof (plugin as any).asUser !== "function") { - return objExports; + return guardPluginApi(plugin, objExports); } - return { + return guardPluginApi(plugin, { ...objExports, /** * @deprecated Use appkit.asUser(req) instead. @@ -230,7 +245,7 @@ export class AppKit { */ asUser: (req: import("express").Request) => (plugin as any).asUser(req).exports() as Record, - }; + }); } static async _createApp< diff --git a/packages/appkit/src/core/plugin-context.ts b/packages/appkit/src/core/plugin-context.ts index 6605d8161..27603bdc1 100644 --- a/packages/appkit/src/core/plugin-context.ts +++ b/packages/appkit/src/core/plugin-context.ts @@ -6,6 +6,10 @@ import { normalizeIdentityError, } from "../context/execution-context"; import { createRequestScope } from "../context/request-scope"; +import { + assertPluginExecution, + getPluginResourceTypes, +} from "../context/resource-capabilities"; import { ServiceContext } from "../context/service-context"; import { createLogger } from "../logging/logger"; import { @@ -242,6 +246,7 @@ export class PluginContext { : timeoutSignal; try { + assertPluginExecution(provider); const result = await provider.executeAgentTool( toolName, args, @@ -268,9 +273,10 @@ export class PluginContext { // Inherit the caller or establish request user scope before the span and tool run. return getCallerContext() ? executeInCurrentScope() - : createRequestScope(req, this.createCallerContext).run( - executeInCurrentScope, - ); + : createRequestScope(req, getPluginResourceTypes(provider), { + createCaller: this.createCallerContext, + legacy: true, + }).run(executeInCurrentScope); } /** diff --git a/packages/appkit/src/plugin/interceptors/telemetry.ts b/packages/appkit/src/plugin/interceptors/telemetry.ts index 4759897ff..2463605d4 100644 --- a/packages/appkit/src/plugin/interceptors/telemetry.ts +++ b/packages/appkit/src/plugin/interceptors/telemetry.ts @@ -1,7 +1,7 @@ import type { TelemetryConfig } from "shared"; import { - getCurrentUserId, + getCurrentPrincipalId, isInUserContext, normalizeIdentityError, } from "../../context/execution-context"; @@ -56,7 +56,7 @@ export class TelemetryInterceptor implements ExecutionInterceptor { "execution.context", isInUserContext() ? "user" : "service", ); - span.setAttribute("caller.id", getCurrentUserId()); + span.setAttribute("caller.id", getCurrentPrincipalId()); if (isDevOboFallback()) { span.setAttribute("execution.obo_dev_fallback", true); } diff --git a/packages/appkit/src/plugin/plugin.ts b/packages/appkit/src/plugin/plugin.ts index 3b0d7418b..7eed23826 100644 --- a/packages/appkit/src/plugin/plugin.ts +++ b/packages/appkit/src/plugin/plugin.ts @@ -16,10 +16,14 @@ import { camelToKebab } from "shared"; import { AppManager } from "../app"; import { CacheManager } from "../cache"; -import { getCurrentUserId } from "../context"; +import { getCurrentPrincipalId } from "../context"; import { warnContextDeprecation } from "../context/deprecation"; import { normalizeIdentityError } from "../context/execution-context"; import { createRequestScope } from "../context/request-scope"; +import { + assertPluginExecution, + getPluginResourceTypes, +} from "../context/resource-capabilities"; import { scopePlugin } from "../context/scoped-api"; import type { PluginContext } from "../core/plugin-context"; import { AppKitError, AuthenticationError } from "../errors"; @@ -344,14 +348,38 @@ export abstract class Plugin< protected resolveUserId(req: express.Request): string { const userId = req.header("x-forwarded-user")?.trim(); if (userId) return userId; - if (process.env.NODE_ENV === "development") return getCurrentUserId(); + if (process.env.NODE_ENV === "development") return getCurrentPrincipalId(); throw AuthenticationError.missingUserId(); } /** @deprecated Use appkit.asUser(req) to scope the whole app. */ asUser(req: express.Request): this { warnContextDeprecation("Plugin.asUser", "appkit.asUser(req)"); - return scopePlugin(this, createRequestScope(req)); + assertPluginExecution(this); + return scopePlugin( + this, + createRequestScope(req, getPluginResourceTypes(this), { legacy: true }), + ); + } + + /** + * User-scoped executor for the plugin's own on-behalf-of routing (for + * example `.obo.sql` lanes), using the same caller context as the + * app-level `appkit.asUser(req)`. + * + * Unlike the deprecated public `asUser(req)`, this does not emit a + * deprecation warning and does not opt out of the app-only resource + * guards, so internal routing matches current on-behalf-of semantics + * rather than the retained legacy behavior. + * + * @internal + */ + protected _asUserScoped(req: express.Request): this { + assertPluginExecution(this); + return scopePlugin( + this, + createRequestScope(req, getPluginResourceTypes(this)), + ); } // streaming execution with interceptors @@ -361,6 +389,7 @@ export abstract class Plugin< options: StreamExecutionSettings, userKey?: string, ) { + assertPluginExecution(this); // destructure options const { stream: streamConfig, @@ -375,7 +404,7 @@ export abstract class Plugin< }); // get user key from context if not provided - const effectiveUserKey = userKey ?? getCurrentUserId(); + const effectiveUserKey = userKey ?? getCurrentPrincipalId(); const self = this; // capture the active OTel context (HTTP span) before entering the async generator, @@ -449,12 +478,13 @@ export abstract class Plugin< options: PluginExecutionSettings, userKey?: string, ): Promise> { + assertPluginExecution(this); const executeConfig = this._buildExecutionConfig(options); const interceptors = this._buildInterceptors(executeConfig); // get user key from context if not provided - const effectiveUserKey = userKey ?? getCurrentUserId(); + const effectiveUserKey = userKey ?? getCurrentPrincipalId(); const context: InterceptorContext = { metadata: new Map(), diff --git a/packages/appkit/src/plugin/tests/plugin.test.ts b/packages/appkit/src/plugin/tests/plugin.test.ts index 64484fb0c..981791e7b 100644 --- a/packages/appkit/src/plugin/tests/plugin.test.ts +++ b/packages/appkit/src/plugin/tests/plugin.test.ts @@ -325,7 +325,7 @@ describe("Plugin", () => { {}, // The plugin forwards the resolved user key as the 4th argument to // bind the stream to its creator. The test passes `false` as an - // explicit override, which propagates through `userKey ?? getCurrentUserId()`. + // explicit override, which propagates through `userKey ?? getCurrentPrincipalId()`. false, ); }); diff --git a/packages/appkit/src/plugins/ai-search/ai-search.ts b/packages/appkit/src/plugins/ai-search/ai-search.ts index 4c62cf7a8..f74260110 100644 --- a/packages/appkit/src/plugins/ai-search/ai-search.ts +++ b/packages/appkit/src/plugins/ai-search/ai-search.ts @@ -8,7 +8,7 @@ import type { VsQueryParams, VsRawResponse, } from "../../connectors/ai-search/types"; -import { getCurrentUserId, getWorkspaceClient } from "../../context"; +import { getCurrentPrincipalId, getWorkspaceClient } from "../../context"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; @@ -170,7 +170,7 @@ export class AiSearchPlugin extends Plugin { // trusted and keep the override.) const { columns: _clientColumns, ...safeBody } = body; const isAsUser = indexConfig.auth === "on-behalf-of-user"; - const plugin = isAsUser ? this.asUser(req) : this; + const plugin = isAsUser ? this._asUserScoped(req) : this; const queryType = safeBody.queryType ?? indexConfig.queryType ?? "hybrid"; @@ -237,7 +237,9 @@ export class AiSearchPlugin extends Plugin { try { const plugin = - indexConfig.auth === "on-behalf-of-user" ? this.asUser(req) : this; + indexConfig.auth === "on-behalf-of-user" + ? this._asUserScoped(req) + : this; // Uncached: a page token is a single-use cursor. `querySettings` has // no `cacheKey`, so caching stays off. @@ -316,7 +318,7 @@ export class AiSearchPlugin extends Plugin { // execute so a cache hit skips the embedding and the VS call. const { queryType } = this._resolveQueryParams(request, indexConfig); - // getCurrentUserId() is the user's id under asUser(), the service id + // getCurrentPrincipalId() is the user's id under asUser(), the service id // otherwise — keying per caller like the route's executorKey. const result = await this.execute( async (signal) => { @@ -327,7 +329,7 @@ export class AiSearchPlugin extends Plugin { signal, ); }, - this._executeSettings(request, indexConfig, getCurrentUserId()), + this._executeSettings(request, indexConfig, getCurrentPrincipalId()), ); if (!result.ok) { diff --git a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts index b692c1d67..f9277932f 100644 --- a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -11,8 +11,8 @@ import { Context } from "../../../workspace-client"; vi.mock("../../../context", () => ({ getWorkspaceClient: vi.fn(() => mockWorkspaceClient), - getCurrentUserId: vi.fn(() => "test-user"), - // OBO plumbing so asUser() runs its non-dev path. getCurrentUserId stays + getCurrentPrincipalId: vi.fn(() => "test-user"), + // OBO plumbing so asUser() runs its non-dev path. getCurrentPrincipalId stays // constant, so per-user scoping is driven by executorKey in the cacheKey. runInCallerContext: (_ctx: unknown, fn: () => T): T => fn(), ServiceContext: { diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index 987ac51d3..4e8bdf077 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -210,7 +210,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { statementId: string, ): Promise { const attempts: Array<() => Promise> = [ - () => this.asUser(req)._getColumnNames(statementId), + () => this._asUserScoped(req)._getColumnNames(statementId), () => this._getColumnNames(statementId), ]; for (const attempt of attempts) { @@ -231,7 +231,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { /** * Fetch column names in the current execution context. Proxied by `asUser`, * so `getWorkspaceClient()` resolves to the user's client when invoked via - * `this.asUser(req)` and the service principal's otherwise. + * `asUser(req)` and the service principal's otherwise. */ async _getColumnNames(statementId: string): Promise { return this.SQLClient.getColumnNames(getWorkspaceClient(), statementId); @@ -311,7 +311,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { } // get execution context - user-scoped if .obo.sql, otherwise service principal - const executor = isAsUser ? this.asUser(req) : this; + const executor = isAsUser ? this._asUserScoped(req) : this; const executorKey = isAsUser ? this.resolveUserId(req) : "global"; const hashedQuery = this.queryProcessor.hashQuery(query); @@ -549,7 +549,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { let executorKey: string; try { const isObo = registration.lane === "obo"; - executor = isObo ? this.asUser(req) : this; + executor = isObo ? this._asUserScoped(req) : this; executorKey = deriveMetricExecutorKey({ lane: registration.lane, userIdentity: isObo ? this.resolveUserId(req) : undefined, @@ -823,7 +823,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { isAsUser: boolean, parameters: IAnalyticsQueryRequest["parameters"], ): Promise { - const executor = isAsUser ? this.asUser(req) : this; + const executor = isAsUser ? this._asUserScoped(req) : this; const executorKey = isAsUser ? this.resolveUserId(req) : "global"; const abortController = new AbortController(); const onClose = () => abortController.abort(); @@ -1051,7 +1051,7 @@ export class AnalyticsPlugin extends Plugin implements ToolProvider { * const result = await analytics.query("SELECT * FROM table") * * // User context execution (in route handler) - * const result = await this.asUser(req).query("SELECT * FROM table") + * const result = await appkit.asUser(req).analytics.query("SELECT * FROM table") * ``` */ async query( diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index a8cbf3692..a25e18daa 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -100,6 +100,10 @@ describe("Analytics Plugin", () => { }); test("/query/:query_key should execute as service principal for .sql files (isAsUser: false)", async () => { + const warnings: string[] = []; + const warn = vi + .spyOn(console, "warn") + .mockImplementation((...args) => warnings.push(args.join(" "))); const plugin = new AnalyticsPlugin(config); const { router, getHandler } = createMockRouter(); @@ -162,6 +166,10 @@ describe("Analytics Plugin", () => { ); expect(mockRes.end).toHaveBeenCalled(); + expect(warnings).not.toContainEqual( + expect.stringContaining("getCurrentUserId is deprecated"), + ); + warn.mockRestore(); }); test("/query/:query_key should execute as user for .obo.sql files (isAsUser: true)", async () => { @@ -827,10 +835,12 @@ describe("Analytics Plugin", () => { // Warehouse readiness must run through the user-context executor too, so // `getWorkspaceClient()` resolves to the user (not the SP) for `.obo.sql`. const ensureReadyMock = vi.fn().mockResolvedValue(undefined); - const asUserSpy = vi.spyOn(plugin as any, "asUser").mockReturnValue({ - query: userExecutorQuery, - _ensureArrowWarehouseReady: ensureReadyMock, - }); + const asUserSpy = vi + .spyOn(plugin as any, "_asUserScoped") + .mockReturnValue({ + query: userExecutorQuery, + _ensureArrowWarehouseReady: ensureReadyMock, + }); const streamExternalLinksMock = vi.fn(function* (_chunks: unknown) { yield new Uint8Array([1, 2, 3]); @@ -980,7 +990,7 @@ describe("Analytics Plugin", () => { .fn() .mockRejectedValue(new Error("RESOURCE_DOES_NOT_EXIST")); const asUserSpy = vi - .spyOn(plugin as any, "asUser") + .spyOn(plugin as any, "_asUserScoped") .mockReturnValue({ _getColumnNames: userGetColumnNames }); plugin.injectRoutes(router); diff --git a/packages/appkit/src/plugins/analytics/tests/metric.test.ts b/packages/appkit/src/plugins/analytics/tests/metric.test.ts index 111ee0a95..abb8952cd 100644 --- a/packages/appkit/src/plugins/analytics/tests/metric.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/metric.test.ts @@ -2954,7 +2954,7 @@ describe("metric route — lane dispatch", () => { ); const { router, getHandler } = createMockRouter(); - const asUserSpy = vi.spyOn(plugin, "asUser"); + const asUserSpy = vi.spyOn(plugin as any, "_asUserScoped"); const executeMock = vi .fn() .mockResolvedValue({ result: { data: [{ arr: 1 }] } }); @@ -2994,7 +2994,7 @@ describe("metric route — lane dispatch", () => { ); const { router, getHandler } = createMockRouter(); - const asUserSpy = vi.spyOn(plugin, "asUser"); + const asUserSpy = vi.spyOn(plugin as any, "_asUserScoped"); const executeMock = vi .fn() .mockResolvedValue({ result: { data: [{ arr: 1 }] } }); diff --git a/packages/appkit/src/plugins/files/plugin.ts b/packages/appkit/src/plugins/files/plugin.ts index 8136055df..1800b6265 100644 --- a/packages/appkit/src/plugins/files/plugin.ts +++ b/packages/appkit/src/plugins/files/plugin.ts @@ -20,8 +20,7 @@ import { } from "../../connectors/files"; import { getCallerContext, - getCurrentActorId, - getCurrentUserId, + getCurrentPrincipalId, getExecutionContext, getWorkspaceClient, runInCallerContext, @@ -178,7 +177,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { "falling back to service principal identity (dev mode). " + "In production this request would 401.", ); - return { id: getCurrentUserId(), isServicePrincipal: true }; + return { id: getCurrentPrincipalId(), isServicePrincipal: true }; } if (!token) { throw AuthenticationError.missingToken( @@ -211,7 +210,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { "OBO volume requested without x-forwarded-access-token — falling back to service principal identity (dev mode). " + "In production this request would 401.", ); - return { id: getCurrentUserId(), isServicePrincipal: true }; + return { id: getCurrentPrincipalId(), isServicePrincipal: true }; } throw AuthenticationError.missingToken( !token @@ -292,7 +291,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { logger.debug( "No x-forwarded-user header — proceeding with service principal identity for policy evaluation.", ); - user = { id: getCurrentUserId(), isServicePrincipal: true }; + user = { id: getCurrentPrincipalId(), isServicePrincipal: true }; } } } catch (error) { @@ -566,7 +565,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { authMode: "service-principal" | "on-behalf-of-user", ): PluginExecutionSettings { // OBO volumes: disable list/read cache. The cache layer is keyed by - // `getCurrentUserId()`, so user A's writes can only invalidate user A's + // `getCurrentPrincipalId()`, so user A's writes can only invalidate user A's // cache entry — user B would continue to see stale data for the same // volume/path until TTL. Disabling caching trades read performance for // correctness; the alternative is a per-(volume, path) generation @@ -633,7 +632,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { * * On OBO volumes the read cache is disabled (see `_readSettings`), so * invalidation is a no-op here for `mode === "on-behalf-of-user"`. The - * cache layer is keyed by `getCurrentUserId()`, so user A's writes can + * cache layer is keyed by `getCurrentPrincipalId()`, so user A's writes can * only invalidate user A's cache entry — user B would otherwise see stale * data for the same volume/path until TTL. Disabling the cache on OBO * trades read performance for correctness; the alternative is a @@ -664,7 +663,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { return; } const parent = parentDirectory(writtenPath); - const userKey = getCurrentUserId(); + const userKey = getCurrentPrincipalId(); const tryDelete = async (segment: string): Promise => { try { await this.cache.delete( @@ -1451,11 +1450,11 @@ export class FilesPlugin extends Plugin implements ToolProvider { /** * Run `fn` under the correct execution context. * - `userCtx` is `null`: invokes `fn` directly so the service-principal - * `WorkspaceClient` and `getCurrentUserId()` are used — identical + * `WorkspaceClient` and `getCurrentPrincipalId()` are used, identical * behavior to pre-OBO releases. This covers both SP volumes and the * OBO dev-fallback path (where headers were missing). * - `userCtx` is a `CallerContext`: wraps `fn` in `runInCallerContext(userCtx)`, - * so SDK calls execute as the end user and `getCurrentUserId()` (and + * so SDK calls execute as the end user and `getCurrentPrincipalId()` (and * therefore cache keys) resolve to the user's ID. * * The caller is responsible for building `userCtx` exactly once per @@ -1565,7 +1564,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { * `runInCallerContext(userCtx, ...)`. Used by `VolumeHandle.asUser(req)` to * force the SDK identity to the end user regardless of the volume's * `auth` setting. The policy check baked into each method (via - * `createVolumeAPI`) runs inside the same scope, so `getCurrentUserId()` + * `createVolumeAPI`) runs inside the same scope, so `getCurrentPrincipalId()` * and any cache `userKey` derived from it also resolve to the user. * * Each wrapped invocation tags the connector's span with @@ -1842,7 +1841,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { // policy must see the user too (not a hardcoded service principal). const policyUser: FilePolicyUser = { get id() { - return getCurrentActorId() ?? ServiceContext.get().serviceUserId; + return getCurrentPrincipalId(); }, get isServicePrincipal() { return getCallerContext() === undefined; diff --git a/packages/appkit/src/plugins/files/tests/plugin.test.ts b/packages/appkit/src/plugins/files/tests/plugin.test.ts index b1e306dec..f9ae1ad01 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.test.ts @@ -3,6 +3,7 @@ import { PassThrough, Readable } from "node:stream"; import { mockServiceContext, setupDatabricksEnv } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { getCurrentPrincipalId } from "../../../context"; import { createRequestScope } from "../../../context/request-scope"; import { scopeApi } from "../../../context/scoped-api"; import { ServiceContext } from "../../../context/service-context"; @@ -63,7 +64,7 @@ vi.mock("../../../context", async (importOriginal) => { return { ...actual, getWorkspaceClient: vi.fn(() => mockClient), - getCurrentUserId: vi.fn(() => "test-service-principal"), + getCurrentPrincipalId: vi.fn(() => "test-service-principal"), }; }); @@ -2272,19 +2273,19 @@ describe("FilesPlugin", () => { } /** - * Replace the default `getCurrentUserId` mock with one that delegates to + * Replace the default `getCurrentPrincipalId` mock with one that delegates to * the real implementation, so that calls inside `runInUserContext` resolve * to the wrapped UserContext's `userId` (and the per-user cache key * derived from it). */ - async function useRealGetCurrentUserId() { + async function useRealGetCurrentPrincipalId() { const actual = await vi.importActual( "../../../context", ); const ctx = await import("../../../context"); - vi.mocked(ctx.getCurrentUserId).mockImplementation( - actual.getCurrentUserId, + vi.mocked(ctx.getCurrentPrincipalId).mockImplementation( + actual.getCurrentPrincipalId, ); } @@ -2305,7 +2306,7 @@ describe("FilesPlugin", () => { }); test("OBO list + valid token wraps SDK call in user context (alice's userId resolves inside the wrapped fn)", async () => { - await useRealGetCurrentUserId(); + await useRealGetCurrentPrincipalId(); const policySpy = vi.fn().mockReturnValue(true); const plugin = new FilesPlugin({ volumes: { @@ -2323,9 +2324,9 @@ describe("FilesPlugin", () => { const observedUserIds: string[] = []; mockClient.files.listDirectoryContents.mockImplementation( async function* () { - // getCurrentUserId() inside the wrapped fn should resolve to alice. + // getCurrentPrincipalId() inside the wrapped fn should resolve to alice. const ctx = await import("../../../context"); - observedUserIds.push(ctx.getCurrentUserId()); + observedUserIds.push(ctx.getCurrentPrincipalId()); yield { name: "o.txt", path: "/o.txt", is_directory: false }; }, ); @@ -2407,14 +2408,14 @@ describe("FilesPlugin", () => { }); test("OBO read cache is DISABLED: cross-user reads do not share cache state", async () => { - // Per Fix 3: the read cache is keyed by `getCurrentUserId()`, so user + // Per Fix 3: the read cache is keyed by `getCurrentPrincipalId()`, so user // A's writes can only invalidate user A's cache entry. Cross-user // staleness was the bug. The chosen mitigation (Option B) is to // disable the read cache on OBO volumes entirely — this test pins // that contract: OBO reads must NOT consult `getOrExecute`. The // alternative (Option A: per-(volume, path) generation counters) // would re-enable cache here. - await useRealGetCurrentUserId(); + await useRealGetCurrentPrincipalId(); const policySpy = vi.fn().mockReturnValue(true); const plugin = new FilesPlugin({ volumes: { @@ -2458,7 +2459,7 @@ describe("FilesPlugin", () => { }); test("SP volume reads still use the cache (cache is only disabled for OBO)", async () => { - await useRealGetCurrentUserId(); + await useRealGetCurrentPrincipalId(); const plugin = new FilesPlugin({ volumes: { obo_vol: { @@ -2635,18 +2636,18 @@ describe("FilesPlugin", () => { } /** - * Replace the default `getCurrentUserId` mock with the real implementation + * Replace the default `getCurrentPrincipalId` mock with the real implementation * so calls inside `runInUserContext` resolve to the wrapped UserContext's * `userId` (mirrors the helper used by the `OBO read routes` block). */ - async function useRealGetCurrentUserId() { + async function useRealGetCurrentPrincipalId() { const actual = await vi.importActual( "../../../context", ); const ctx = await import("../../../context"); - vi.mocked(ctx.getCurrentUserId).mockImplementation( - actual.getCurrentUserId, + vi.mocked(ctx.getCurrentPrincipalId).mockImplementation( + actual.getCurrentPrincipalId, ); } @@ -2706,7 +2707,7 @@ describe("FilesPlugin", () => { * this assertion fails. SP-token would silently leak to UC otherwise. */ test("OBO upload: outgoing fetch PUT carries user-token Authorization header (not SP)", async () => { - await useRealGetCurrentUserId(); + await useRealGetCurrentPrincipalId(); await useRealGetWorkspaceClient(); // SP-token marker — what the existing mockClient would inject if the @@ -2892,7 +2893,7 @@ describe("FilesPlugin", () => { }); test("OBO delete + valid token + UC denies → user-token client invoked, error propagated", async () => { - await useRealGetCurrentUserId(); + await useRealGetCurrentPrincipalId(); await useRealGetWorkspaceClient(); // Distinct user-token client with a `files.delete` that mimics a UC @@ -3091,14 +3092,14 @@ describe("FilesPlugin", () => { */ test("SP write awaits cache.delete BEFORE sending the response (no write→read race)", async () => { // Restore the default (mocked) `getWorkspaceClient` and - // `getCurrentUserId`, since earlier tests in this block install the + // `getCurrentPrincipalId`, since earlier tests in this block install the // REAL implementations and `vi.clearAllMocks` does NOT reset // implementations. const ctx = await import("../../../context"); vi.mocked(ctx.getWorkspaceClient).mockImplementation( () => mockClient as any, ); - vi.mocked(ctx.getCurrentUserId).mockImplementation( + vi.mocked(ctx.getCurrentPrincipalId).mockImplementation( () => "test-service-principal", ); @@ -3170,22 +3171,22 @@ describe("FilesPlugin", () => { /** * Fix 3 regression: cross-user OBO read freshness. The OBO read cache - * is keyed by `getCurrentUserId()`, so user A's writes can only + * is keyed by `getCurrentPrincipalId()`, so user A's writes can only * invalidate user A's cache entry. With cache disabled on OBO, user B * must see fresh data after user A writes. */ test("OBO write by user A → user B's next read sees fresh data (cross-user freshness; cache disabled on OBO)", async () => { // This test relies on the DEFAULT mocked `getWorkspaceClient` and - // `getCurrentUserId` (always returning the SP fixture's + // `getCurrentPrincipalId` (always returning the SP fixture's // `mockClient`). Earlier tests in this block install the REAL impls - // via `useRealGetWorkspaceClient`/`useRealGetCurrentUserId`, and + // via `useRealGetWorkspaceClient`/`useRealGetCurrentPrincipalId`, and // Vitest's `vi.clearAllMocks` between tests does NOT reset // implementations — so we restore the defaults explicitly. const ctx = await import("../../../context"); vi.mocked(ctx.getWorkspaceClient).mockImplementation( () => mockClient as any, ); - vi.mocked(ctx.getCurrentUserId).mockImplementation( + vi.mocked(ctx.getCurrentPrincipalId).mockImplementation( () => "test-service-principal", ); @@ -3845,6 +3846,15 @@ describe("FilesPlugin", () => { ); test("policy user reflects the active scope on the default surface", async () => { + // This file pins getCurrentPrincipalId to the SP id; use the real one so + // the default surface reports the caller under appkit.asUser(req). + const ctx = + await vi.importActual( + "../../../context", + ); + vi.mocked(getCurrentPrincipalId).mockImplementation( + ctx.getCurrentPrincipalId, + ); const seen: FilePolicyUser[] = []; const recorder = (_a: unknown, _r: unknown, user: FilePolicyUser) => { seen.push({ id: user.id, isServicePrincipal: user.isServicePrincipal }); @@ -3858,13 +3868,20 @@ describe("FilesPlugin", () => { ); const handle = plugin.exports()("uploads"); - await handle.list("d"); - await appkitAsUser(handle).list("d"); - await handle - .asUser({ - header: (n: string) => (oboHeaders as any)[n.toLowerCase()], - } as any) - .list("d"); + try { + await handle.list("d"); + await appkitAsUser(handle).list("d"); + await handle + .asUser({ + header: (n: string) => (oboHeaders as any)[n.toLowerCase()], + } as any) + .list("d"); + } finally { + // Restore the file-wide pin so later tests are unaffected. + vi.mocked(getCurrentPrincipalId).mockImplementation( + () => "test-service-principal", + ); + } expect(seen[0].isServicePrincipal).toBe(true); expect(seen[1]).toEqual({ diff --git a/packages/appkit/src/plugins/genie/genie.ts b/packages/appkit/src/plugins/genie/genie.ts index 351d3988e..4e51c2c2a 100644 --- a/packages/appkit/src/plugins/genie/genie.ts +++ b/packages/appkit/src/plugins/genie/genie.ts @@ -136,7 +136,7 @@ export class GeniePlugin extends Plugin implements ToolProvider { method: "post", path: "/:alias/messages", handler: async (req: express.Request, res: express.Response) => { - await this.asUser(req)._handleSendMessage(req, res); + await this._asUserScoped(req)._handleSendMessage(req, res); }, }); @@ -145,7 +145,7 @@ export class GeniePlugin extends Plugin implements ToolProvider { method: "get", path: "/:alias/conversations/:conversationId", handler: async (req: express.Request, res: express.Response) => { - await this.asUser(req)._handleGetConversation(req, res); + await this._asUserScoped(req)._handleGetConversation(req, res); }, }); @@ -154,7 +154,7 @@ export class GeniePlugin extends Plugin implements ToolProvider { method: "get", path: "/:alias/conversations/:conversationId/messages/:messageId", handler: async (req: express.Request, res: express.Response) => { - await this.asUser(req)._handleGetMessage(req, res); + await this._asUserScoped(req)._handleGetMessage(req, res); }, }); } diff --git a/packages/appkit/src/plugins/jobs/plugin.ts b/packages/appkit/src/plugins/jobs/plugin.ts index fe2af9599..d9f4913b1 100644 --- a/packages/appkit/src/plugins/jobs/plugin.ts +++ b/packages/appkit/src/plugins/jobs/plugin.ts @@ -9,7 +9,7 @@ import type { import { toJSONSchema } from "zod"; import { JobsConnector } from "../../connectors/jobs"; -import { getCurrentUserId, getWorkspaceClient } from "../../context"; +import { getCurrentPrincipalId, getWorkspaceClient } from "../../context"; import { ExecutionError, ValidationError } from "../../errors"; import { createLogger } from "../../logging/logger"; import type { ExecutionResult } from "../../plugin"; @@ -258,7 +258,7 @@ class JobsPlugin extends Plugin { const self = this; // Capture client and userId eagerly: the closures below run later, after the ALS context may have exited. const client = this.client; - const userKey = getCurrentUserId(); + const userKey = getCurrentPrincipalId(); /** * Verify that `runId` belongs to this job's configured `jobId`. Returns diff --git a/packages/appkit/src/plugins/lakebase/lakebase.ts b/packages/appkit/src/plugins/lakebase/lakebase.ts index 11e7889d6..5f351d827 100644 --- a/packages/appkit/src/plugins/lakebase/lakebase.ts +++ b/packages/appkit/src/plugins/lakebase/lakebase.ts @@ -14,6 +14,7 @@ import { } from "../../connectors/lakebase"; import { getClientOptions } from "../../context/client-options"; import { getCallerContext } from "../../context/execution-context"; +import { assertResourceExecution } from "../../context/resource-capabilities"; import { buildToolkitEntries } from "../../core/agent/build-toolkit"; import { defineTool, @@ -43,11 +44,8 @@ const OBO_POOL_DEFAULTS = { * Wraps `@databricks/lakebase` to provide a standard `pg.Pool` with automatic * OAuth token refresh, integrated with AppKit's logger and OpenTelemetry setup. * - * Supports On-Behalf-Of (OBO) via `asUser(req)` — each user gets a separate - * `pg.Pool` authenticated with their Databricks identity, enabling features - * like Row-Level Security (RLS). Routing is handled transparently by - * {@link RoutingPool}, which reads the execution context set by the base - * class `asUser()`. + * App-only in v1. Caller-scoped operations reject before opening a connection. + * The platform's postgres user scope is deferred for this connector. * * @example * ```ts @@ -60,8 +58,6 @@ const OBO_POOL_DEFAULTS = { * // Service principal query * const result = await AppKit.lakebase.query("SELECT * FROM users WHERE id = $1", [userId]); * - * // User-scoped query (per-user pool, RLS enforced) - * const mine = await AppKit.lakebase.asUser(req).query("SELECT * FROM my_data"); * ``` */ export class LakebasePlugin extends Plugin implements ToolProvider { @@ -124,8 +120,8 @@ export class LakebasePlugin extends Plugin implements ToolProvider { /** * Executes a parameterized SQL query against the Lakebase pool. * - * When called inside `asUser(req)`, the query automatically routes to - * the per-user pool via {@link RoutingPool}. + * New caller scopes enforce the app-only resource contract. Deprecated + * user entry points retain their established per-user pool routing. * * @param text - SQL query string, using `$1`, `$2`, ... placeholders * @param values - Parameter values corresponding to placeholders @@ -296,9 +292,10 @@ export class LakebasePlugin extends Plugin implements ToolProvider { /** * Returns the pool config for the current execution context. - * Inside `asUser(req)`, returns user-scoped config; otherwise SP config. + * Caller scopes reject; plain execution returns SP configuration. */ private activePoolConfig() { + assertResourceExecution("postgres"); const ctx = getCallerContext(); if (ctx) { const user = ctx.principal.userEmail ?? ctx.principal.userId; @@ -314,14 +311,12 @@ export class LakebasePlugin extends Plugin implements ToolProvider { /** * Returns the plugin's public API, accessible via `AppKit.lakebase`. * - * - `pool` — The connection pool (routes to per-user pool when inside `asUser(req)`) - * - `query` — Convenience method for executing parameterized SQL queries - * - `getOrmConfig()` — Returns a config object compatible with Drizzle, TypeORM, Sequelize, etc. - * Inside `asUser(req)`, returns user-scoped config. - * - `getPgConfig()` — Returns a `pg.PoolConfig` object for manual pool construction. - * Inside `asUser(req)`, returns user-scoped config. + * - `pool`: The service-principal connection pool. + * - `query`: Convenience method for executing parameterized SQL queries. + * - `getOrmConfig()`: SP configuration for Drizzle, TypeORM, Sequelize, etc. + * - `getPgConfig()`: A service-principal `pg.PoolConfig` for manual pools. * - * Use `AppKit.lakebase.asUser(req)` to get the same API backed by a per-user pool. + * Caller-scoped access is rejected in v1, including cached pool handles. */ exports() { return { diff --git a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts index 89aaab9c7..b33eaef5c 100644 --- a/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts +++ b/packages/appkit/src/plugins/lakebase/tests/lakebase-agent-tool.test.ts @@ -68,7 +68,7 @@ function makePlugin( return new LakebasePlugin(config); } -describe("LakebasePlugin — agent tool opt-in", () => { +describe("LakebasePlugin - agent tool opt-in", () => { test("does not register an agent tool by default", () => { const plugin = makePlugin({}); expect(plugin.getAgentTools()).toEqual([]); @@ -106,7 +106,7 @@ describe("LakebasePlugin — agent tool opt-in", () => { }); }); -describe("LakebasePlugin — readOnly enforcement", () => { +describe("LakebasePlugin - readOnly enforcement", () => { let plugin: LakebasePlugin; beforeEach(async () => { @@ -203,7 +203,7 @@ describe("LakebasePlugin — readOnly enforcement", () => { }); }); -describe("LakebasePlugin — shutdown", () => { +describe("LakebasePlugin - shutdown", () => { test("closes the SP pool and all OBO pools via shutdown()", async () => { const { createLakebasePool, createLakebasePoolManager } = await import("../../../connectors/lakebase"); @@ -242,13 +242,13 @@ describe("LakebasePlugin — shutdown", () => { plugin.abortActiveOperations(); // Other plugins' shutdown() hooks may still need database connections - // to drain state — the pools must survive the abort phase. + // to drain state - the pools must survive the abort phase. expect(spPool.end).not.toHaveBeenCalled(); expect(oboManager.closeAll).not.toHaveBeenCalled(); }); }); -describe("LakebasePlugin — destructive mode", () => { +describe("LakebasePlugin - destructive mode", () => { test("does NOT wrap in read-only transaction when readOnly: false", async () => { const queryMock = vi.fn((_text: string, _values?: unknown[]) => Promise.resolve({ rows: [] }), @@ -274,7 +274,7 @@ describe("LakebasePlugin — destructive mode", () => { }); }); -describe("LakebasePlugin — OBO via RoutingPool", () => { +describe("LakebasePlugin - OBO via RoutingPool", () => { const userPoolQueries: Array<{ text: string; values?: unknown[] }> = []; const userClientQueries: Array<{ text: string; values?: unknown[] }> = []; diff --git a/packages/appkit/src/plugins/serving/serving.ts b/packages/appkit/src/plugins/serving/serving.ts index d3c3b8ac2..a0ac141a4 100644 --- a/packages/appkit/src/plugins/serving/serving.ts +++ b/packages/appkit/src/plugins/serving/serving.ts @@ -149,7 +149,7 @@ export class ServingPlugin extends Plugin { method: "post", path: "/:alias/invoke", handler: async (req: express.Request, res: express.Response) => { - await this.asUser(req)._handleInvoke(req, res); + await this._asUserScoped(req)._handleInvoke(req, res); }, }); @@ -158,7 +158,7 @@ export class ServingPlugin extends Plugin { method: "post", path: "/:alias/stream", handler: async (req: express.Request, res: express.Response) => { - await this.asUser(req)._handleStream(req, res); + await this._asUserScoped(req)._handleStream(req, res); }, }); } else { @@ -169,14 +169,14 @@ export class ServingPlugin extends Plugin { res: express.Response, ) => { req.params.alias ??= "default"; - await this.asUser(req)._handleInvoke(req, res); + await this._asUserScoped(req)._handleInvoke(req, res); }; const streamHandler = async ( req: express.Request, res: express.Response, ) => { req.params.alias ??= "default"; - await this.asUser(req)._handleStream(req, res); + await this._asUserScoped(req)._handleStream(req, res); }; this.route(router, { @@ -341,7 +341,7 @@ export class ServingPlugin extends Plugin { return { ...spApi, asUser: (req: express.Request) => { - const userPlugin = this.asUser(req) as ServingPlugin; + const userPlugin = this._asUserScoped(req) as ServingPlugin; return userPlugin.createEndpointAPI(resolved); }, }; diff --git a/packages/appkit/src/plugins/tests/core-obo-no-deprecation.test.ts b/packages/appkit/src/plugins/tests/core-obo-no-deprecation.test.ts new file mode 100644 index 000000000..2b8753d66 --- /dev/null +++ b/packages/appkit/src/plugins/tests/core-obo-no-deprecation.test.ts @@ -0,0 +1,79 @@ +import { + createMockRequest, + createMockResponse, + createMockRouter, + setupDatabricksEnv, + useServiceContextMock, +} from "@tools/test-helpers"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { ServiceContext } from "../../context/service-context"; +import { GeniePlugin } from "../genie/genie"; +import { ServingPlugin } from "../serving/serving"; + +// Regression for the deprecated-asUser warning on core-plugin OBO routes. +// genie and serving route handlers must scope via the internal +// `_asUserScoped(req)`, never the public `Plugin.asUser(req)`, so exercising +// an OBO route emits no "Plugin.asUser is deprecated" warning. This is the +// not-emitted complement to the emitted-case assertion in +// core/tests/appkit-user-scope.test.ts. It lives in its own file so the +// deprecation warn-once dedup starts fresh and cannot mask a regression. +describe("core plugin OBO routes do not warn about deprecated asUser", () => { + useServiceContextMock(); + + let warn: ReturnType; + beforeEach(() => { + setupDatabricksEnv(); + ServiceContext.reset(); + warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + }); + afterEach(() => { + vi.restoreAllMocks(); + ServiceContext.reset(); + }); + + const oboReq = (params: Record) => + createMockRequest({ + params, + body: { content: "hi", input: "hi" }, + headers: { + "x-forwarded-access-token": "user-token", + "x-forwarded-user": "user-1", + }, + }); + + const deprecationWarnings = () => + warn.mock.calls.filter((args) => + args.some( + (a) => + typeof a === "string" && a.includes("Plugin.asUser is deprecated"), + ), + ); + + test("genie sendMessage route scopes without the deprecation warning", async () => { + const plugin = new GeniePlugin({ spaces: { myspace: "space-123" } }); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const handler = getHandler("POST", "/:alias/messages"); + + // Unknown alias returns 404 inside the handler, but only after the route + // wrapper has already scoped via _asUserScoped(req), which is the line + // under test. + await handler(oboReq({ alias: "unknown" }), createMockResponse()); + + expect(deprecationWarnings()).toHaveLength(0); + }); + + test("serving invoke route scopes without the deprecation warning", async () => { + const plugin = new ServingPlugin({ + endpoints: { llm: { env: "DATABRICKS_SERVING_ENDPOINT_NAME" } }, + }); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const handler = getHandler("POST", "/:alias/invoke"); + + await handler(oboReq({ alias: "unknown" }), createMockResponse()); + + expect(deprecationWarnings()).toHaveLength(0); + }); +}); diff --git a/packages/appkit/src/telemetry/tests/telemetry-interceptor.test.ts b/packages/appkit/src/telemetry/tests/telemetry-interceptor.test.ts index fb4327ea3..6a33905ef 100644 --- a/packages/appkit/src/telemetry/tests/telemetry-interceptor.test.ts +++ b/packages/appkit/src/telemetry/tests/telemetry-interceptor.test.ts @@ -40,7 +40,9 @@ describe("TelemetryInterceptor", () => { userKey: "test", }; - vi.spyOn(executionContext, "getCurrentUserId").mockReturnValue("test-user"); + vi.spyOn(executionContext, "getCurrentPrincipalId").mockReturnValue( + "test-user", + ); }); test("should execute function and set span status to OK on success", async () => { @@ -139,7 +141,9 @@ describe("TelemetryInterceptor", () => { test("should set execution context as 'service' when not in user context", async () => { vi.spyOn(executionContext, "isInUserContext").mockReturnValue(false); - vi.spyOn(executionContext, "getCurrentUserId").mockReturnValue("sp-123"); + vi.spyOn(executionContext, "getCurrentPrincipalId").mockReturnValue( + "sp-123", + ); const interceptor = new TelemetryInterceptor(mockTelemetry); const fn = vi.fn().mockResolvedValue("result"); @@ -154,7 +158,9 @@ describe("TelemetryInterceptor", () => { test("should set execution context as 'user' when in user context", async () => { vi.spyOn(executionContext, "isInUserContext").mockReturnValue(true); - vi.spyOn(executionContext, "getCurrentUserId").mockReturnValue("user-123"); + vi.spyOn(executionContext, "getCurrentPrincipalId").mockReturnValue( + "user-123", + ); const interceptor = new TelemetryInterceptor(mockTelemetry); const fn = vi.fn().mockResolvedValue("result"); diff --git a/packages/appkit/src/testing/test-plugin-context.ts b/packages/appkit/src/testing/test-plugin-context.ts index 97a433224..af8fa7e00 100644 --- a/packages/appkit/src/testing/test-plugin-context.ts +++ b/packages/appkit/src/testing/test-plugin-context.ts @@ -325,10 +325,10 @@ function createTestPluginContextSync( // Reuse production header validation and ALS. Only the client is faked. const asUser = (req: IAppRequest): ToolProvider => { record.asUserRequests.push(req as express.Request); - const scope = createRequestScope( - req as express.Request, - createCallerContext, - ); + const scope = createRequestScope(req as express.Request, [], { + createCaller: createCallerContext, + legacy: true, + }); return { ...base, diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index 42abeb903..7641302b4 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -10,7 +10,11 @@ export * from "./genie"; export * from "./metric-filter"; export * from "./metric-metadata"; export * from "./plugin"; -export { pluginManifestSchema } from "./schemas/manifest"; +export { + pluginManifestSchema, + APP_ONLY_RESOURCE_TYPES, + SCOPE_BY_TYPE, +} from "./schemas/manifest"; export * from "./sql"; export * from "./sse/analytics"; export * from "./tunnel";