diff --git a/docs/docs/api/appkit/Class.Plugin.md b/docs/docs/api/appkit/Class.Plugin.md index a5002a097..b4f9e5221 100644 --- a/docs/docs/api/appkit/Class.Plugin.md +++ b/docs/docs/api/appkit/Class.Plugin.md @@ -555,7 +555,7 @@ Resolve the effective user ID from a request. Returns the `x-forwarded-user` header when present. In development mode (`NODE_ENV=development`) falls back to the current context user ID so -that callers outside an active `runInUserContext` scope still get a +that callers outside an active `runInCallerContext` scope still get a consistent value. #### Parameters diff --git a/docs/docs/api/appkit/Function.getCurrentActorId.md b/docs/docs/api/appkit/Function.getCurrentActorId.md new file mode 100644 index 000000000..688cf9ae9 --- /dev/null +++ b/docs/docs/api/appkit/Function.getCurrentActorId.md @@ -0,0 +1,11 @@ +# Function: getCurrentActorId() + +```ts +function getCurrentActorId(): string | undefined; +``` + +The initiating user in a caller scope; no user actor exists in service scope. + +## Returns + +`string` \| `undefined` diff --git a/docs/docs/api/appkit/Function.getCurrentPrincipalKey.md b/docs/docs/api/appkit/Function.getCurrentPrincipalKey.md new file mode 100644 index 000000000..7c790388e --- /dev/null +++ b/docs/docs/api/appkit/Function.getCurrentPrincipalKey.md @@ -0,0 +1,11 @@ +# Function: getCurrentPrincipalKey() + +```ts +function getCurrentPrincipalKey(): string; +``` + +Get the principal key for future cache keying: `app` or `user:`. + +## Returns + +`string` diff --git a/docs/docs/api/appkit/Function.getExecutionContext.md b/docs/docs/api/appkit/Function.getExecutionContext.md index a89ab4133..14751649d 100644 --- a/docs/docs/api/appkit/Function.getExecutionContext.md +++ b/docs/docs/api/appkit/Function.getExecutionContext.md @@ -1,17 +1,20 @@ # Function: getExecutionContext() ```ts -function getExecutionContext(): ExecutionContext; +function getExecutionContext(): + | ServiceContextState + | CallerContext & UserContext; ``` Get the current execution context. -- If running inside a user context (via asUser), returns the user context +- If running inside a caller context (via asUser), returns the caller context - Otherwise, returns the service context ## Returns -`ExecutionContext` + \| `ServiceContextState` + \| [`CallerContext`](Interface.CallerContext.md) & `UserContext` ## Throws diff --git a/docs/docs/api/appkit/Function.getWarehouseId.md b/docs/docs/api/appkit/Function.getWarehouseId.md new file mode 100644 index 000000000..bc3ae1c00 --- /dev/null +++ b/docs/docs/api/appkit/Function.getWarehouseId.md @@ -0,0 +1,21 @@ +# Function: getWarehouseId() + +```ts +function getWarehouseId(): Promise; +``` + +Get the configured SQL warehouse ID after app initialization. +The warehouse is an app resource; SP and caller executions use the same binding. +Deprecated user-context scopes retain support for explicit warehouse overrides. + +## Returns + +`Promise`\<`string`\> + +## Throws + +ConfigurationError if no SQL warehouse was required at startup. + +## Throws + +InitializationError if the app resources are not initialized. diff --git a/docs/docs/api/appkit/Interface.CallerContext.md b/docs/docs/api/appkit/Interface.CallerContext.md new file mode 100644 index 000000000..f9890e18a --- /dev/null +++ b/docs/docs/api/appkit/Interface.CallerContext.md @@ -0,0 +1,37 @@ +# Interface: CallerContext + +Caller identity and workspace for one immutable execution scope. + +## Properties + +### client + +```ts +readonly client: WorkspaceClient; +``` + +*** + +### principal + +```ts +readonly principal: CallerPrincipal; +``` + +*** + +### tokenFingerprint? + +```ts +readonly optional tokenFingerprint: string; +``` + +Truncated SHA-256 hash of the caller token, used to detect rotation. + +*** + +### workspaceId + +```ts +readonly workspaceId: Promise; +``` diff --git a/docs/docs/api/appkit/Interface.PluginManifest.md b/docs/docs/api/appkit/Interface.PluginManifest.md index d0db2c533..32762abeb 100644 --- a/docs/docs/api/appkit/Interface.PluginManifest.md +++ b/docs/docs/api/appkit/Interface.PluginManifest.md @@ -56,6 +56,22 @@ schema: JSONSchema7; *** +### deprecated? + +```ts +optional deprecated: boolean; +``` + +When true, the plugin is deprecated. It still ships and functions, but tooling (e.g. `appkit plugin list`) may hide or flag it. The recommended replacement is noted in the plugin description. + +#### Inherited from + +```ts +Omit.deprecated +``` + +*** + ### description ```ts @@ -277,6 +293,27 @@ Omit.scaffolding *** +### scopes? + +```ts +optional scopes: ( + | "ai-gateway" + | "mcp.external" + | "mcp.functions" + | "workspace.workspace" + | "catalog.catalogs:read" + | "catalog.schemas:read" + | "catalog.tables:read")[]; +``` + +#### Inherited from + +```ts +Omit.scopes +``` + +*** + ### stability? ```ts diff --git a/docs/docs/api/appkit/TypeAlias.CallerPrincipal.md b/docs/docs/api/appkit/TypeAlias.CallerPrincipal.md new file mode 100644 index 000000000..53a9c82bb --- /dev/null +++ b/docs/docs/api/appkit/TypeAlias.CallerPrincipal.md @@ -0,0 +1,12 @@ +# Type Alias: CallerPrincipal + +```ts +type CallerPrincipal = Readonly<{ + type: "user"; + userEmail?: string; + userId: string; + userName?: string; +}>; +``` + +The caller identity whose permissions authorize execution, not its resources. diff --git a/docs/docs/api/appkit/TypeAlias.ExecutionContext.md b/docs/docs/api/appkit/TypeAlias.ExecutionContext.md new file mode 100644 index 000000000..4cf921385 --- /dev/null +++ b/docs/docs/api/appkit/TypeAlias.ExecutionContext.md @@ -0,0 +1,5 @@ +# Type Alias: ExecutionContext + +```ts +type ExecutionContext = ServiceContextState | CallerContext; +``` diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index 36e5e6d58..7f5668c39 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -48,6 +48,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [AutoInheritToolsConfig](Interface.AutoInheritToolsConfig.md) | Auto-inherit configuration. When enabled for a given agent origin, agents with no explicit `tools:` declaration receive every registered ToolProvider plugin tool whose author marked `autoInheritable: true`. Tools without that flag — destructive, state-mutating, or privilege-sensitive — never spread automatically and must be wired via `tools:` (object or function form in code, `plugin:NAME` entries in markdown frontmatter). | | [BasePluginConfig](Interface.BasePluginConfig.md) | Base configuration interface for AppKit plugins | | [CacheConfig](Interface.CacheConfig.md) | Configuration for the CacheInterceptor. Controls TTL, size limits, storage backend, and probabilistic cleanup. | +| [CallerContext](Interface.CallerContext.md) | Caller identity and workspace for one immutable execution scope. | | [CustomJudgeSpec](Interface.CustomJudgeSpec.md) | A custom LLM-judge definition: a prompt template and choice→score mapping. | | [DatabaseCredential](Interface.DatabaseCredential.md) | Database credentials with OAuth token for Postgres connection | | [DatabaseRegistry](Interface.DatabaseRegistry.md) | CANONICAL augmentation target. Empty by default; the generated `database.d.ts` augments it via `declare module "@databricks/appkit" { interface DatabaseRegistry { ... } }`. | @@ -140,6 +141,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [AgentTools](TypeAlias.AgentTools.md) | Per-agent tool record. String keys map to inline tools, toolkit entries, hosted tools, etc. | | [AgentToolsFn](TypeAlias.AgentToolsFn.md) | Function form of `AgentDefinition.tools`. Receives the typed [Plugins](TypeAlias.Plugins.md) map and returns a tool record. Invoked exactly once at setup (or once per `runAgent` call in standalone mode); the result is cached as the agent's resolved tool record. | | [BaseSystemPromptOption](TypeAlias.BaseSystemPromptOption.md) | - | +| [CallerPrincipal](TypeAlias.CallerPrincipal.md) | The caller identity whose permissions authorize execution, not its resources. | | [ConfigSchema](TypeAlias.ConfigSchema.md) | Configuration schema definition for plugin config. Re-exported from the standard JSON Schema Draft 7 types. | | [DatabaseApiConfig](TypeAlias.DatabaseApiConfig.md) | Full generated CRUD for every declared table by default. Set false to disable all generated routes, or use an object to restrict tables and writes. Keyed routes require a public primary key; upsert stays programmatic. Route names must start with a letter, contain only letters, digits, `_`, or `-`, be at most 64 characters, and be unique ignoring case. Invalid names fail setup; exclude internal tables with `api.tables` or use `api: false`. | | [DatabaseApiWriteOperation](TypeAlias.DatabaseApiWriteOperation.md) | Generated HTTP write operations. | @@ -147,6 +149,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [DatabaseExports](TypeAlias.DatabaseExports.md) | Typed database API published by the plugin. | | [EntityHooks](TypeAlias.EntityHooks.md) | Response shaping and mutation lifecycle declared for one table. | | [EvalProgress](TypeAlias.EvalProgress.md) | - | +| [ExecutionContext](TypeAlias.ExecutionContext.md) | - | | [ExecutionResult](TypeAlias.ExecutionResult.md) | Discriminated union for plugin execution results. | | [FileAction](TypeAlias.FileAction.md) | Every action the files plugin can perform. | | [FilePolicy](TypeAlias.FilePolicy.md) | A policy function that decides whether `user` may perform `action` on `resource`. Return `true` to allow, `false` to deny. | @@ -224,12 +227,15 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [fromSupervisorApi](Function.fromSupervisorApi.md) | Creates an [AgentAdapter](Interface.AgentAdapter.md) backed by the Databricks AI Gateway Responses API (`/ai-gateway/mlflow/v1/responses`). | | [functionToolToDefinition](Function.functionToolToDefinition.md) | - | | [generateDatabaseCredential](Function.generateDatabaseCredential.md) | Generate OAuth credentials for Postgres database connection using the proper Postgres API. | +| [getCurrentActorId](Function.getCurrentActorId.md) | The initiating user in a caller scope; no user actor exists in service scope. | +| [getCurrentPrincipalKey](Function.getCurrentPrincipalKey.md) | Get the principal key for future cache keying: `app` or `user:`. | | [getExecutionContext](Function.getExecutionContext.md) | Get the current execution context. | | [getLakebaseOrmConfig](Function.getLakebaseOrmConfig.md) | Get Lakebase connection configuration for ORMs that don't accept pg.Pool directly. | | [getLakebasePgConfig](Function.getLakebasePgConfig.md) | Get Lakebase connection configuration for PostgreSQL clients. | | [getPluginManifest](Function.getPluginManifest.md) | Loads and validates the manifest from a plugin constructor. Normalizes string type/permission to strict ResourceType/ResourcePermission. | | [getResourceRequirements](Function.getResourceRequirements.md) | Gets the resource requirements from a plugin's manifest. | | [getUsernameWithApiLookup](Function.getUsernameWithApiLookup.md) | Resolves the PostgreSQL username for a Lakebase connection. | +| [getWarehouseId](Function.getWarehouseId.md) | Get the configured SQL warehouse ID after app initialization. The warehouse is an app resource; SP and caller executions use the same binding. Deprecated user-context scopes retain support for explicit warehouse overrides. | | [getWorkspaceClient](Function.getWorkspaceClient.md) | Get workspace client from config or SDK default auth chain | | [id](Function.id.md) | - | | [includes](Function.includes.md) | Passes when the value contains `substring`. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 4414a2b34..58e5b00c7 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -172,6 +172,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.CacheConfig", label: "CacheConfig" }, + { + type: "doc", + id: "api/appkit/Interface.CallerContext", + label: "CallerContext" + }, { type: "doc", id: "api/appkit/Interface.CustomJudgeSpec", @@ -613,6 +618,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/TypeAlias.BaseSystemPromptOption", label: "BaseSystemPromptOption" }, + { + type: "doc", + id: "api/appkit/TypeAlias.CallerPrincipal", + label: "CallerPrincipal" + }, { type: "doc", id: "api/appkit/TypeAlias.ConfigSchema", @@ -648,6 +658,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/TypeAlias.EvalProgress", label: "EvalProgress" }, + { + type: "doc", + id: "api/appkit/TypeAlias.ExecutionContext", + label: "ExecutionContext" + }, { type: "doc", id: "api/appkit/TypeAlias.ExecutionResult", @@ -995,6 +1010,16 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Function.generateDatabaseCredential", label: "generateDatabaseCredential" }, + { + type: "doc", + id: "api/appkit/Function.getCurrentActorId", + label: "getCurrentActorId" + }, + { + type: "doc", + id: "api/appkit/Function.getCurrentPrincipalKey", + label: "getCurrentPrincipalKey" + }, { type: "doc", id: "api/appkit/Function.getExecutionContext", @@ -1025,6 +1050,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Function.getUsernameWithApiLookup", label: "getUsernameWithApiLookup" }, + { + type: "doc", + id: "api/appkit/Function.getWarehouseId", + label: "getWarehouseId" + }, { type: "doc", id: "api/appkit/Function.getWorkspaceClient", diff --git a/docs/docs/plugins/analytics.md b/docs/docs/plugins/analytics.md index be969696d..b144a5c08 100644 --- a/docs/docs/plugins/analytics.md +++ b/docs/docs/plugins/analytics.md @@ -36,6 +36,22 @@ await createApp({ The execution context is determined by the SQL file name, not by the hook call. +### Warehouse binding + +The warehouse is an app resource, not part of the caller's identity. After +`createApp()` initializes Analytics, custom plugins can read the binding with +`getWarehouseId` from `@databricks/appkit`. It returns a `Promise` and +uses the same app warehouse for service-principal and caller executions. +Credentials still come from the current execution context. + +`ServiceContextState.warehouseId` and the old context-module `getWarehouseId` +helper are deprecated compatibility accessors. Use the package-root helper for +new code. Warehouse environment lookup and development discovery are unchanged. + +The service-context property is retained for one release with a one-time +deprecation warning, then scheduled for removal. Internal callers use the +resource accessor already. + ## SQL parameters Use `:paramName` placeholders and optionally annotate parameter types using SQL comments: diff --git a/docs/static/schemas/plugin-manifest.schema.json b/docs/static/schemas/plugin-manifest.schema.json index 236cd8706..fd726d41a 100644 --- a/docs/static/schemas/plugin-manifest.schema.json +++ b/docs/static/schemas/plugin-manifest.schema.json @@ -4,6 +4,22 @@ "title": "AppKit Plugin Manifest", "type": "object", "properties": { + "scopes": { + "description": "Capability-only user_api_scopes with no resource ID.", + "type": "array", + "items": { + "type": "string", + "enum": [ + "ai-gateway", + "mcp.external", + "mcp.functions", + "workspace.workspace", + "catalog.catalogs:read", + "catalog.schemas:read", + "catalog.tables:read" + ] + } + }, "$schema": { "description": "Reference to the JSON Schema for validation", "type": "string" diff --git a/packages/appkit/src/connectors/lakebase/routing-pool.ts b/packages/appkit/src/connectors/lakebase/routing-pool.ts index c62a08957..1a8ad41d6 100644 --- a/packages/appkit/src/connectors/lakebase/routing-pool.ts +++ b/packages/appkit/src/connectors/lakebase/routing-pool.ts @@ -1,7 +1,7 @@ import type { Pool, PoolClient, QueryResult, QueryResultRow } from "pg"; -import { getUserContext } from "../../context/execution-context"; -import type { UserContext } from "../../context/user-context"; +import type { CallerContext } from "../../context/caller-context"; +import { getCallerContext } from "../../context/execution-context"; /** * Subset of `pg.Pool` exposed by the Lakebase plugin. @@ -26,7 +26,7 @@ export interface LakebasePool { * A `pg.Pool`-like wrapper that routes queries to the appropriate pool * based on the current execution context. * - * When called inside `runInUserContext()` (set up by `Plugin.asUser(req)`), + * 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. * @@ -37,11 +37,11 @@ export interface LakebasePool { export class RoutingPool implements LakebasePool { constructor( private spPool: Pool, - private resolveUserPool: (ctx: UserContext) => Pool, + private resolveUserPool: (ctx: CallerContext) => Pool, ) {} private activePool(): Pool { - const userCtx = getUserContext(); + 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 4ccccd86d..29912b098 100644 --- a/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts +++ b/packages/appkit/src/connectors/lakebase/tests/routing-pool.test.ts @@ -52,7 +52,13 @@ describe("RoutingPool", () => { expect(result.rows).toEqual([{ source: "user" }]); expect(userPool.query).toHaveBeenCalledWith("SELECT 1", undefined); expect(spPool.query).not.toHaveBeenCalled(); - expect(resolveUserPool).toHaveBeenCalledWith(userCtx); + expect(resolveUserPool).toHaveBeenCalledWith( + expect.objectContaining({ + client: userCtx.client, + workspaceId: userCtx.workspaceId, + principal: expect.objectContaining({ type: "user", userId: "user-1" }), + }), + ); }); test("connect() routes to user pool inside runInUserContext", async () => { diff --git a/packages/appkit/src/context/caller-context.ts b/packages/appkit/src/context/caller-context.ts new file mode 100644 index 000000000..49500a8bd --- /dev/null +++ b/packages/appkit/src/context/caller-context.ts @@ -0,0 +1,39 @@ +import type { ServiceContextState } from "./service-context"; + +/** The caller identity whose permissions authorize execution, not its resources. */ +export type CallerPrincipal = Readonly<{ + type: "user"; + userId: string; + userName?: string; + userEmail?: string; +}>; + +/** Caller identity and workspace for one immutable execution scope. */ +export interface CallerContext { + readonly client: ServiceContextState["client"]; + readonly principal: CallerPrincipal; + /** Truncated SHA-256 hash of the caller token, used to detect rotation. */ + readonly tokenFingerprint?: string; + readonly workspaceId: Promise; +} + +const snapshots = new WeakSet(); + +/** Snapshot identity without freezing the SDK client's internal lifecycle. */ +export function snapshotCallerContext(ctx: CallerContext): CallerContext { + if (snapshots.has(ctx)) return ctx; + const snapshot = Object.freeze({ + client: ctx.client, + principal: Object.freeze({ ...ctx.principal }), + tokenFingerprint: ctx.tokenFingerprint, + workspaceId: ctx.workspaceId, + }); + snapshots.add(snapshot); + return snapshot; +} + +export type ExecutionContext = ServiceContextState | CallerContext; + +export function isCallerContext(ctx: ExecutionContext): ctx is CallerContext { + return "principal" in ctx && ctx.principal.type === "user"; +} diff --git a/packages/appkit/src/context/deprecation.ts b/packages/appkit/src/context/deprecation.ts new file mode 100644 index 000000000..30d26bf16 --- /dev/null +++ b/packages/appkit/src/context/deprecation.ts @@ -0,0 +1,13 @@ +import { createLogger } from "../logging/logger"; + +const logger = createLogger("execution-context"); +const warned = new Set(); + +export function warnContextDeprecation( + name: string, + replacement: string, +): void { + if (warned.has(name)) return; + warned.add(name); + logger.warn(`${name} is deprecated. Use ${replacement} instead.`); +} diff --git a/packages/appkit/src/context/execution-context.ts b/packages/appkit/src/context/execution-context.ts index 1acd569a5..dc5ea6bd4 100644 --- a/packages/appkit/src/context/execution-context.ts +++ b/packages/appkit/src/context/execution-context.ts @@ -1,86 +1,130 @@ import { AsyncLocalStorage } from "node:async_hooks"; -import { ConfigurationError } from "../errors"; -import { ServiceContext } from "./service-context"; import { - type ExecutionContext, - isUserContext, + captureWarehouseId, + getWarehouseId as getResourceWarehouseId, + runWithResourceBindings, + type WarehouseBinding, +} from "../resources/warehouse"; +import { type CallerContext, snapshotCallerContext } from "./caller-context"; +import { warnContextDeprecation } from "./deprecation"; +import { ServiceContext, type ServiceContextState } from "./service-context"; +import { + legacyUserContext, + toCallerContext, type UserContext, } from "./user-context"; -/** - * AsyncLocalStorage for execution context. - * Used to pass user context through the call stack without explicit parameters. - */ -const executionContextStorage = new AsyncLocalStorage(); +const executionContextStorage = new AsyncLocalStorage(); + +function runInCallerScope( + callerContext: CallerContext, + fn: () => T, + legacyResources?: WarehouseBinding, +): T { + const caller = snapshotCallerContext(callerContext); + return runWithResourceBindings(legacyResources, () => + executionContextStorage.run(caller, fn), + ); +} /** - * Run a function in the context of a user. - * All calls within the function will have access to the user context. + * Run a function with an immutable snapshot of the caller context. + * Nested and concurrent scopes keep their own identities. * - * @param userContext - The user context to use + * @param callerContext - The caller context to use * @param fn - The function to run * @returns The result of the function */ -export function runInUserContext(userContext: UserContext, fn: () => T): T { - return executionContextStorage.run(userContext, fn); +export function runInCallerContext( + callerContext: CallerContext, + fn: () => T, +): T { + return runInCallerScope(callerContext, fn); +} + +/** @deprecated Use runInCallerContext. */ +export function runInUserContext( + userContext: UserContext | (CallerContext & Pick), + fn: () => T, +): T { + warnContextDeprecation("runInUserContext", "runInCallerContext"); + if (!("principal" in userContext) || "warehouseId" in userContext) { + return runInCallerScope( + toCallerContext(userContext), + fn, + Object.freeze({ warehouseId: userContext.warehouseId }), + ); + } + return runInCallerContext(toCallerContext(userContext), fn); } /** * Get the current execution context. * - * - If running inside a user context (via asUser), returns the user context + * - If running inside a caller context (via asUser), returns the caller context * - Otherwise, returns the service context * * @throws Error if ServiceContext is not initialized */ -export function getExecutionContext(): ExecutionContext { - const userContext = executionContextStorage.getStore(); - if (userContext) { - return userContext; +export function getExecutionContext(): + | ServiceContextState + | (CallerContext & UserContext) { + const callerContext = executionContextStorage.getStore(); + if (callerContext) { + return legacyUserContext(callerContext, captureWarehouseId()); } return ServiceContext.get(); } /** - * Get the current user ID for cache keying and telemetry. - * - * Returns the user ID if in user context, otherwise the service user ID. + * Get the principal key for future cache keying: `app` or `user:`. + */ +export function getCurrentPrincipalKey(): string { + const caller = getCallerContext(); + return caller ? `user:${caller.principal.userId}` : "app"; +} + +/** The initiating user in a caller scope; no user actor exists in service scope. */ +export function getCurrentActorId(): string | undefined { + return getCallerContext()?.principal.userId; +} + +/** + * @deprecated Use getCurrentPrincipalKey for new cache keys or getCurrentActorId + * for audit. Preserves the bare user or service ID for existing callers. */ export function getCurrentUserId(): string { - const ctx = getExecutionContext(); - if (isUserContext(ctx)) { - return ctx.userId; - } - return ctx.serviceUserId; + warnContextDeprecation( + "getCurrentUserId", + "getCurrentPrincipalKey (cache) or getCurrentActorId (audit)", + ); + return getCurrentActorId() ?? ServiceContext.get().serviceUserId; } /** * Get the WorkspaceClient for the current execution context. */ export function getWorkspaceClient() { - return getExecutionContext().client; + return (getCallerContext() ?? ServiceContext.get()).client; } /** - * Get the warehouse ID promise. + * @deprecated Import getWarehouseId from @databricks/appkit instead of context. */ export function getWarehouseId(): Promise { - const ctx = getExecutionContext(); - if (!ctx.warehouseId) { - throw ConfigurationError.resourceNotFound( - "Warehouse ID", - "No plugin requires a SQL Warehouse. Add a sql_warehouse resource to your plugin manifest, or set DATABRICKS_WAREHOUSE_ID", - ); - } - return ctx.warehouseId; + warnContextDeprecation( + "context.getWarehouseId", + "getWarehouseId() from @databricks/appkit", + ); + return getResourceWarehouseId(); } /** * Get the workspace ID promise. */ export function getWorkspaceId(): Promise { - return getExecutionContext().workspaceId; + return (getCallerContext() ?? ServiceContext.get()).workspaceId; } /** @@ -92,10 +136,17 @@ export function isInUserContext(): boolean { } /** - * Get the user context if one is active, otherwise `undefined`. + * Get the caller context if one is active, otherwise `undefined`. * Unlike `getExecutionContext()`, this does not require `ServiceContext` * to be initialized and never throws. */ -export function getUserContext(): UserContext | undefined { +export function getCallerContext(): CallerContext | undefined { return executionContextStorage.getStore(); } + +/** @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; +} diff --git a/packages/appkit/src/context/index.ts b/packages/appkit/src/context/index.ts index d306d359e..63a319ef8 100644 --- a/packages/appkit/src/context/index.ts +++ b/packages/appkit/src/context/index.ts @@ -1,10 +1,20 @@ export { + getCallerContext, + getCurrentActorId, + getCurrentPrincipalKey, getCurrentUserId, getExecutionContext, getWarehouseId, getWorkspaceClient, getWorkspaceId, runInUserContext, + runInCallerContext, } from "./execution-context"; +export { + type CallerContext, + type CallerPrincipal, + type ExecutionContext, + isCallerContext, +} from "./caller-context"; export { ServiceContext } from "./service-context"; export type { UserContext } from "./user-context"; diff --git a/packages/appkit/src/context/service-context.ts b/packages/appkit/src/context/service-context.ts index 5d9ad78fe..0475e8a95 100644 --- a/packages/appkit/src/context/service-context.ts +++ b/packages/appkit/src/context/service-context.ts @@ -5,34 +5,39 @@ import { ConfigurationError, InitializationError, } from "../errors"; +import { WarehouseResource } from "../resources/warehouse"; import { type ClientOptions, ConfigError, createWorkspaceClient, - type sql, type WorkspaceClient, } from "../workspace-client"; +import { type CallerContext, snapshotCallerContext } from "./caller-context"; import { getClientOptions } from "./client-options"; -import type { UserContext } from "./user-context"; +import { warnContextDeprecation } from "./deprecation"; +import { legacyUserContext, type UserContext } from "./user-context"; /** - * Service context holds the service principal client and shared resources. + * Service context holds the service principal identity and workspace client. * This is initialized once at app startup and shared across all requests. */ export interface ServiceContextState { /** WorkspaceClient authenticated as the service principal */ - client: WorkspaceClient; + readonly client: WorkspaceClient; /** The service principal's user ID */ - serviceUserId: string; - /** Promise that resolves to the warehouse ID (only present when a plugin requires `SQL_WAREHOUSE` resource) */ - warehouseId?: Promise; + readonly serviceUserId: string; + /** + * @deprecated Use getWarehouseId() from @databricks/appkit. + * Retained for backward compatibility. + */ + readonly warehouseId?: Promise; /** Promise that resolves to the workspace ID */ - workspaceId: Promise; + readonly workspaceId: Promise; } /** * ServiceContext is a singleton that manages the service principal's - * WorkspaceClient and shared resources like warehouse/workspace IDs. + * WorkspaceClient and workspace ID. WarehouseResource owns warehouse bindings. * * It's initialized once at app startup and provides the foundation * for both service principal and user context execution. @@ -45,12 +50,12 @@ export class ServiceContext { * Initialize the service context. Should be called once at app startup. * Safe to call multiple times - will return the same instance. * - * @param options - Which shared resources to resolve (derived from plugin manifests). + * @param options - A resolved warehouse ID, or a boolean enabling discovery. * @param client - Optional pre-configured WorkspaceClient to use instead * of creating one from environment credentials. */ static async initialize( - options?: { warehouseId?: boolean }, + options?: { warehouseId?: boolean | string }, client?: WorkspaceClient, ): Promise { if (ServiceContext.instance) { @@ -88,19 +93,20 @@ export class ServiceContext { } /** - * Create a user context from request headers. + * Create an immutable caller context from the existing user request headers. * * @param token - The user's access token from x-forwarded-access-token header * @param userId - The user's ID from x-forwarded-user header * @param userName - Optional user name + * @param userEmail - Optional email from x-forwarded-email * @throws Error if token is not provided */ - static createUserContext( + static createCallerContext( token: string, userId: string, userName?: string, userEmail?: string, - ): UserContext { + ): CallerContext { if (!token) { throw AuthenticationError.missingToken("user token"); } @@ -127,16 +133,33 @@ export class ServiceContext { .digest("hex") .slice(0, 16); - return { + return snapshotCallerContext({ client: userClient, + principal: { type: "user", userId, userName, userEmail }, + tokenFingerprint, + workspaceId: serviceCtx.workspaceId, + }); + } + + /** @deprecated Use ServiceContext.createCallerContext. */ + static createUserContext( + token: string, + userId: string, + userName?: string, + userEmail?: string, + ): CallerContext & UserContext { + warnContextDeprecation( + "ServiceContext.createUserContext", + "ServiceContext.createCallerContext", + ); + const caller = ServiceContext.createCallerContext( + token, userId, userName, userEmail, - tokenFingerprint, - warehouseId: serviceCtx.warehouseId, - workspaceId: serviceCtx.workspaceId, - isUserContext: true, - }; + ); + const warehouseId = WarehouseResource.get()?.warehouseId; + return legacyUserContext(caller, () => warehouseId); } /** @@ -148,37 +171,38 @@ export class ServiceContext { } private static async createContext( - options?: { warehouseId?: boolean }, + options?: { warehouseId?: boolean | string }, client?: WorkspaceClient, ): Promise { try { const wsClient = client ?? createWorkspaceClient({ clientOptions: getClientOptions() }); - const [resolvedWorkspaceId, currentUser, resolvedWarehouseId] = + const [resolvedWorkspaceId, currentUser, resolvedResources] = await Promise.all([ ServiceContext.getWorkspaceId(wsClient), wsClient.currentUser.me(), - options?.warehouseId - ? ServiceContext.getWarehouseId(wsClient) - : Promise.resolve(undefined as string | undefined), + WarehouseResource.resolve(wsClient, options?.warehouseId), ]); if (!currentUser.id) { throw ConfigurationError.resourceNotFound("Service user ID"); } - const warehouseId = - options?.warehouseId && resolvedWarehouseId !== undefined - ? Promise.resolve(resolvedWarehouseId) - : undefined; + const resources = WarehouseResource.bind(resolvedResources); - return { + return Object.freeze({ client: wsClient, serviceUserId: currentUser.id, - warehouseId, + get warehouseId() { + warnContextDeprecation( + "ServiceContextState.warehouseId", + "getWarehouseId() from @databricks/appkit", + ); + return resources.warehouseId; + }, workspaceId: Promise.resolve(resolvedWorkspaceId), - }; + }); } catch (e) { if (e instanceof ConfigError) { throw ConfigurationError.databricksAuthenticationSetupFailed( @@ -213,76 +237,12 @@ export class ServiceContext { return response["x-databricks-org-id"]; } - private static async getWarehouseId( - client: WorkspaceClient, - ): Promise { - if (process.env.DATABRICKS_WAREHOUSE_ID) { - return process.env.DATABRICKS_WAREHOUSE_ID; - } - - const agenticMode = - process.env.DATABRICKS_APPS_AGENTIC_MODE === "true" || - process.env.DATABRICKS_APPS_AGENTIC_MODE === "1"; - - if (process.env.NODE_ENV === "development" && !agenticMode) { - const response = (await client.apiClient.request({ - path: "/api/2.0/sql/warehouses", - method: "GET", - headers: new Headers(), - raw: false, - query: { - skip_cannot_use: "true", - }, - })) as { warehouses: sql.EndpointInfo[] }; - - const priorities: Record = { - RUNNING: 0, - STOPPED: 1, - STARTING: 2, - STOPPING: 3, - DELETED: 99, - DELETING: 99, - }; - - const warehouses = (response.warehouses || []).sort((a, b) => { - return ( - priorities[a.state as sql.State] - priorities[b.state as sql.State] - ); - }); - - if (response.warehouses.length === 0) { - throw ConfigurationError.resourceNotFound( - "Warehouse ID", - "Please configure the DATABRICKS_WAREHOUSE_ID environment variable", - ); - } - - const firstWarehouse = warehouses[0]; - if ( - firstWarehouse.state === "DELETED" || - firstWarehouse.state === "DELETING" || - !firstWarehouse.id - ) { - throw ConfigurationError.resourceNotFound( - "Warehouse ID", - "Please configure the DATABRICKS_WAREHOUSE_ID environment variable", - ); - } - - return firstWarehouse.id; - } - - throw ConfigurationError.resourceNotFound( - "Warehouse ID", - "Please configure the DATABRICKS_WAREHOUSE_ID environment variable", - ); - } - /** * Reset the service context. Only for testing purposes. */ static reset(): void { ServiceContext.instance = null; ServiceContext.initPromise = null; + WarehouseResource.reset(); } } diff --git a/packages/appkit/src/context/tests/execution-context.test.ts b/packages/appkit/src/context/tests/execution-context.test.ts new file mode 100644 index 000000000..9668335ef --- /dev/null +++ b/packages/appkit/src/context/tests/execution-context.test.ts @@ -0,0 +1,393 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { getWarehouseId } from "../../resources"; +import { WarehouseResource } from "../../resources/warehouse"; +import { createMockWorkspaceClient } from "../../testing/mock-workspace-client"; +import * as workspaceClient from "../../workspace-client"; +import { getUserContext } from "../execution-context"; +import { + type CallerContext, + type UserContext, + getCallerContext, + getCurrentActorId, + getCurrentPrincipalKey, + getCurrentUserId, + getExecutionContext, + getWarehouseId as getLegacyWarehouseId, + getWorkspaceClient, + isCallerContext, + runInCallerContext, + runInUserContext, + ServiceContext, +} from "../index"; +import { + type ExecutionContext, + immutableCallerContext, + isUserContext, +} from "../user-context"; + +const { warn } = vi.hoisted(() => ({ warn: vi.fn() })); +vi.mock("../../logging/logger", () => ({ + createLogger: () => ({ warn }), +})); + +const service = Object.freeze({ + client: createMockWorkspaceClient(), + serviceUserId: "service-123", + workspaceId: Promise.resolve("workspace-123"), +} satisfies ExecutionContext); + +const appResources = Object.freeze({ + warehouseId: Promise.resolve("warehouse-123"), +}); + +function caller(userId: string): CallerContext { + return { + client: createMockWorkspaceClient(), + principal: { type: "user", userId }, + workspaceId: service.workspaceId, + }; +} + +describe("caller execution context", () => { + beforeEach(() => { + vi.spyOn(ServiceContext, "get").mockReturnValue(service); + vi.spyOn(WarehouseResource, "get").mockReturnValue(appResources); + }); + + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllEnvs(); + }); + + test("defaults to the SP and has no initiating user", () => { + expect(getExecutionContext()).toBe(service); + expect(getWorkspaceClient()).toBe(service.client); + expect(getCurrentPrincipalKey()).toBe("app"); + expect(getCurrentActorId()).toBeUndefined(); + expect(getCallerContext()).toBeUndefined(); + expect(isCallerContext(service)).toBe(false); + }); + + test("caller accessors work without an initialized service context", () => { + vi.mocked(ServiceContext.get).mockImplementation(() => { + throw new Error("not initialized"); + }); + expect(getCallerContext()).toBeUndefined(); + expect(getCurrentActorId()).toBeUndefined(); + runInCallerContext(caller("alice"), () => { + expect(getCurrentPrincipalKey()).toBe("user:alice"); + expect(getCurrentActorId()).toBe("alice"); + }); + }); + + test("isolates concurrent callers across asynchronous work", async () => { + let release!: () => void; + const barrier = new Promise((resolve) => { + release = resolve; + }); + const alice = runInCallerContext(caller("alice"), async () => { + await barrier; + expect(getCurrentActorId()).toBe("alice"); + return getCurrentPrincipalKey(); + }); + const bob = runInCallerContext(caller("bob"), async () => { + release(); + await Promise.resolve(); + expect(getCurrentActorId()).toBe("bob"); + return getCurrentPrincipalKey(); + }); + expect(await Promise.all([alice, bob])).toEqual(["user:alice", "user:bob"]); + expect(getExecutionContext()).toBe(service); + }); + + test("restores the parent after a nested scope throws or rejects", async () => { + await runInCallerContext(caller("alice"), async () => { + expect(() => + runInCallerContext(caller("bob"), () => { + throw new Error("sync failure"); + }), + ).toThrow("sync failure"); + expect(getCurrentActorId()).toBe("alice"); + await expect( + runInCallerContext(caller("bob"), async () => { + await Promise.resolve(); + throw new Error("async failure"); + }), + ).rejects.toThrow("async failure"); + expect(getCurrentActorId()).toBe("alice"); + }); + expect(getExecutionContext()).toBe(service); + }); + + test("snapshots the input and freezes context and principal, but not the client", () => { + const input = { + ...caller("alice"), + principal: { type: "user" as const, userId: "alice" }, + }; + runInCallerContext(input, () => { + input.principal.userId = "bob"; + const active = getCallerContext(); + if (!active) throw new Error("Expected an active caller"); + expect(getCurrentPrincipalKey()).toBe("user:alice"); + expect(Object.isFrozen(active)).toBe(true); + expect(active).not.toHaveProperty("userId"); + expect(active).not.toHaveProperty("userName"); + expect(active).not.toHaveProperty("userEmail"); + expect(active).not.toHaveProperty("isUserContext"); + expect(Object.isFrozen(active.principal)).toBe(true); + expect(Reflect.set(active.principal, "userId", "charlie")).toBe(false); + expect(Reflect.set(active, "principal", input.principal)).toBe(false); + expect(active.client).toBe(input.client); + expect(Object.isFrozen(active.client)).toBe(false); + }); + }); + + test("creates the same token client and immutable user principal", () => { + vi.stubEnv("DATABRICKS_HOST", "https://workspace.example.com"); + const client = createMockWorkspaceClient(); + const createClient = vi + .spyOn(workspaceClient, "createWorkspaceClient") + .mockReturnValue(client); + const context = ServiceContext.createCallerContext( + "token", + "alice", + "Alice", + "alice@example.com", + ); + expect(context.principal).toEqual({ + type: "user", + userId: "alice", + userName: "Alice", + userEmail: "alice@example.com", + }); + expect(createClient).toHaveBeenCalledWith( + expect.objectContaining({ + token: "token", + host: "https://workspace.example.com", + authType: "pat", + }), + ); + expect(context.client).toBe(client); + expect(context).not.toHaveProperty("warehouseId"); + expect(context.workspaceId).toBe(service.workspaceId); + expect(context.tokenFingerprint).toHaveLength(16); + expect(Object.isFrozen(context)).toBe(true); + expect(Object.isFrozen(context.principal)).toBe(true); + }); + + test("legacy aliases preserve IDs and fields and warn once per name", () => { + vi.stubEnv("DATABRICKS_HOST", "https://workspace.example.com"); + vi.spyOn(workspaceClient, "createWorkspaceClient").mockReturnValue( + service.client, + ); + warn.mockClear(); + for (let i = 0; i < 2; i++) { + const context = ServiceContext.createUserContext( + "token", + "alice", + "Alice", + "alice@example.com", + ); + expect( + runInUserContext(context, () => { + expect(getCurrentUserId()).toBe("alice"); + expect(getCurrentPrincipalKey()).toBe("user:alice"); + expect(isUserContext(getExecutionContext())).toBe(true); + const legacy = getUserContext(); + if (!legacy) throw new Error("Expected an active legacy context"); + expect(legacy.userId).toBe("alice"); + expect(legacy.userName).toBe("Alice"); + expect(legacy.userEmail).toBe("alice@example.com"); + expect(legacy.isUserContext).toBe(true); + expect(legacy.warehouseId).toBe(appResources.warehouseId); + expect(getLegacyWarehouseId()).toBe(getWarehouseId()); + return 42; + }), + ).toBe(42); + expect(getCurrentUserId()).toBe("service-123"); + } + for (const name of [ + "ServiceContext.createUserContext", + "runInUserContext", + "getCurrentUserId", + "getUserContext", + "isUserContext", + "UserContext.userId", + "UserContext.userName", + "UserContext.userEmail", + "UserContext.isUserContext", + "UserContext.warehouseId", + "context.getWarehouseId", + ]) { + expect( + warn.mock.calls.filter(([message]) => + message.startsWith(`${name} is deprecated.`), + ), + ).toHaveLength(1); + } + }); + + test("accepts an existing legacy context without mutating it", () => { + const legacy: UserContext = { + client: service.client, + workspaceId: service.workspaceId, + userId: "alice", + isUserContext: true, + }; + runInUserContext(legacy, () => { + expect(getCurrentActorId()).toBe("alice"); + expect(getCurrentPrincipalKey()).toBe("user:alice"); + }); + expect(legacy).not.toHaveProperty("principal"); + expect(Object.isFrozen(legacy)).toBe(false); + }); + + test("shares the app warehouse without changing the caller's client", async () => { + expect(getWarehouseId()).toBe(appResources.warehouseId); + const alice = caller("alice"); + await runInCallerContext(alice, async () => { + await Promise.resolve(); + expect(getCallerContext()).not.toHaveProperty("warehouseId"); + expect(getExecutionContext().warehouseId).toBe(appResources.warehouseId); + expect(getWarehouseId()).toBe(appResources.warehouseId); + expect(await getWarehouseId()).toBe("warehouse-123"); + expect(getWorkspaceClient()).toBe(alice.client); + }); + expect(getWorkspaceClient()).toBe(service.client); + }); + + test("reports a missing app warehouse in both service and caller scopes", () => { + vi.mocked(WarehouseResource.get).mockReturnValue({}); + expect(getWarehouseId).toThrow("No plugin requires a SQL Warehouse"); + runInCallerContext(caller("alice"), () => { + expect(getWarehouseId).toThrow("No plugin requires a SQL Warehouse"); + }); + }); + + test("ignores resource fields on canonical caller inputs", () => { + const input = { + ...caller("alice"), + warehouseId: Promise.resolve("not-the-app-warehouse"), + }; + runInCallerContext(input, () => { + expect(getCallerContext()).not.toHaveProperty("warehouseId"); + expect(getWarehouseId()).toBe(appResources.warehouseId); + }); + }); + + test("isolates and snapshots legacy warehouse overrides outside caller identity", async () => { + const warehouseId = Promise.resolve("legacy-warehouse"); + const legacy: UserContext = { + client: createMockWorkspaceClient(), + workspaceId: service.workspaceId, + userId: "alice", + isUserContext: true, + warehouseId, + }; + await Promise.all([ + runInUserContext(legacy, async () => { + legacy.warehouseId = Promise.resolve("changed"); + await Promise.resolve(); + expect(getWarehouseId()).toBe(warehouseId); + expect(getUserContext()?.warehouseId).toBe(warehouseId); + expect(getCallerContext()).not.toHaveProperty("warehouseId"); + expect(getWorkspaceClient()).toBe(legacy.client); + await expect( + runInCallerContext(caller("bob"), async () => { + expect(getWarehouseId()).toBe(appResources.warehouseId); + throw new Error("nested failure"); + }), + ).rejects.toThrow("nested failure"); + expect(getWarehouseId()).toBe(warehouseId); + }), + runInCallerContext(caller("charlie"), async () => { + await Promise.resolve(); + expect(getWarehouseId()).toBe(appResources.warehouseId); + }), + ]); + expect(getWarehouseId()).toBe(appResources.warehouseId); + }); + + test("preserves a missing warehouse in legacy contexts", () => { + runInUserContext( + { + client: service.client, + workspaceId: service.workspaceId, + userId: "alice", + isUserContext: true, + }, + () => { + expect(getUserContext()?.warehouseId).toBeUndefined(); + expect(getWarehouseId).toThrow("No plugin requires a SQL Warehouse"); + }, + ); + }); + + test("public compatibility views keep identity and warehouse snapshots after the scope closes", async () => { + const warehouseId = Promise.resolve("legacy-warehouse"); + const legacy: UserContext = { + client: service.client, + workspaceId: service.workspaceId, + userId: "alice", + isUserContext: true, + warehouseId, + }; + const [execution, user] = runInUserContext( + legacy, + () => [getExecutionContext(), getUserContext()] as const, + ); + legacy.warehouseId = Promise.resolve("changed"); + vi.mocked(WarehouseResource.get).mockReturnValue({}); + await runInCallerContext(caller("bob"), async () => { + await Promise.resolve(); + if (!isUserContext(execution)) throw new Error("Expected user context"); + expect(execution.userId).toBe("alice"); + expect(execution.warehouseId).toBe(warehouseId); + expect(user?.warehouseId).toBe(warehouseId); + expect(Object.isFrozen(execution)).toBe(true); + }); + }); + + test("captured compatibility views retain the app warehouse across rebinding", () => { + const [execution, user] = runInCallerContext( + caller("alice"), + () => [getExecutionContext(), getUserContext()] as const, + ); + vi.mocked(WarehouseResource.get).mockReturnValue({ + warehouseId: Promise.resolve("new-lifecycle"), + }); + expect(execution.warehouseId).toBe(appResources.warehouseId); + expect(user?.warehouseId).toBe(appResources.warehouseId); + }); + + test("uninitialized resources do not prevent reading caller identity", () => { + vi.mocked(WarehouseResource.get).mockReturnValue(undefined); + runInCallerContext(caller("alice"), () => { + const execution = getExecutionContext(); + expect(getCurrentActorId()).toBe("alice"); + expect(getWorkspaceClient()).toBe(execution.client); + expect(() => execution.warehouseId).toThrow("not initialized"); + }); + }); + + test("legacy entry points accept warehouse-free caller contexts", () => { + runInUserContext(caller("alice"), () => { + expect(getWarehouseId()).toBe(appResources.warehouseId); + expect(getUserContext()?.warehouseId).toBe(appResources.warehouseId); + }); + }); + + test("retains the deprecated snapshot helper with legacy identity access", () => { + const original = caller("alice"); + const snapshot = immutableCallerContext(original); + expect(snapshot.userId).toBe("alice"); + expect(snapshot.isUserContext).toBe(true); + expect(Object.isFrozen(snapshot)).toBe(true); + expect(original).not.toHaveProperty("userId"); + runInCallerContext(snapshot, () => { + expect(getCallerContext()).not.toHaveProperty("userId"); + expect(getCurrentActorId()).toBe("alice"); + }); + }); +}); diff --git a/packages/appkit/src/context/tests/service-context.test.ts b/packages/appkit/src/context/tests/service-context.test.ts index f70f85efa..470132142 100644 --- a/packages/appkit/src/context/tests/service-context.test.ts +++ b/packages/appkit/src/context/tests/service-context.test.ts @@ -6,6 +6,8 @@ import { ConfigurationError, InitializationError, } from "../../errors"; +import { getWarehouseId } from "../../resources"; +import { WarehouseResource } from "../../resources/warehouse"; import { ServiceContext } from "../service-context"; // ── Mock the workspace-client wrapper ────────────────────────────── @@ -76,6 +78,8 @@ describe("ServiceContext", () => { expect(state.client).toBe(client); expect(state.serviceUserId).toBe("service-user-123"); expect(await state.workspaceId).toBe("ws-456"); + expect(Object.isFrozen(state)).toBe(true); + expect(Reflect.set(state, "serviceUserId", "another-user")).toBe(false); }); test("should create a WorkspaceClient when none is provided", async () => { @@ -85,25 +89,50 @@ describe("ServiceContext", () => { expect(MockWorkspaceClient).toHaveBeenCalled(); }); - test("should resolve warehouseId when options.warehouseId is true", async () => { + test("the deprecated warehouse getter forwards the binding and warns once", async () => { process.env.DATABRICKS_WAREHOUSE_ID = "wh-789"; - - const state = await ServiceContext.initialize({ warehouseId: true }); - - expect(state.warehouseId).toBeDefined(); - expect(await state.warehouseId).toBe("wh-789"); + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + + try { + const state = await ServiceContext.initialize({ warehouseId: true }); + + expect(await getWarehouseId()).toBe("wh-789"); + expect(warn).not.toHaveBeenCalled(); + expect(state.warehouseId).toBeDefined(); + expect(await state.warehouseId).toBe("wh-789"); + expect(getWarehouseId()).toBe(state.warehouseId); + expect( + Object.getOwnPropertyDescriptor(state, "warehouseId")?.get, + ).toBeTypeOf("function"); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn).toHaveBeenCalledWith( + expect.any(String), + "ServiceContextState.warehouseId is deprecated. Use getWarehouseId() from @databricks/appkit instead.", + ); + } finally { + warn.mockRestore(); + } }); test("should not set warehouseId when options.warehouseId is false", async () => { - const state = await ServiceContext.initialize({ warehouseId: false }); + await ServiceContext.initialize({ warehouseId: false }); + + expect(getWarehouseId).toThrow("No plugin requires a SQL Warehouse"); + }); - expect(state.warehouseId).toBeUndefined(); + test("publishes an explicitly resolved warehouse without discovering another", async () => { + process.env.DATABRICKS_WAREHOUSE_ID = "unrelated-warehouse"; + await ServiceContext.initialize({ warehouseId: "manifest-warehouse" }); + expect(await getWarehouseId()).toBe("manifest-warehouse"); + expect(mockApiRequest).not.toHaveBeenCalledWith( + expect.objectContaining({ path: "/api/2.0/sql/warehouses" }), + ); }); test("should not set warehouseId when options are omitted", async () => { - const state = await ServiceContext.initialize(); + await ServiceContext.initialize(); - expect(state.warehouseId).toBeUndefined(); + expect(getWarehouseId).toThrow("No plugin requires a SQL Warehouse"); }); test("should throw when currentUser.me() returns no id", async () => { @@ -114,6 +143,37 @@ describe("ServiceContext", () => { ); }); + test("does not publish resources when service identity initialization fails", async () => { + mockMe.mockResolvedValue({}); + await expect( + ServiceContext.initialize({ warehouseId: true }), + ).rejects.toThrow(ConfigurationError); + expect(WarehouseResource.get()).toBeUndefined(); + }); + + test("shares one warehouse discovery across concurrent initialization calls", async () => { + delete process.env.DATABRICKS_WAREHOUSE_ID; + process.env.NODE_ENV = "development"; + mockApiRequest.mockImplementation(({ path }: { path: string }) => + Promise.resolve( + path === "/api/2.0/sql/warehouses" + ? { warehouses: [{ id: "wh-shared", state: "RUNNING" }] } + : { "x-databricks-org-id": "ws-456" }, + ), + ); + const [first, second] = await Promise.all([ + ServiceContext.initialize({ warehouseId: true }), + ServiceContext.initialize({ warehouseId: true }), + ]); + expect(first).toBe(second); + expect(await getWarehouseId()).toBe("wh-shared"); + expect( + mockApiRequest.mock.calls.filter( + ([request]) => request.path === "/api/2.0/sql/warehouses", + ), + ).toHaveLength(1); + }); + test("should be idempotent - calling twice returns same instance", async () => { const state1 = await ServiceContext.initialize(); const state2 = await ServiceContext.initialize(); @@ -192,7 +252,7 @@ describe("ServiceContext", () => { expect(userCtx.client).toBeDefined(); }); - test("should share warehouseId and workspaceId from service context", async () => { + test("legacy callers share app resources and the service workspace", async () => { process.env.DATABRICKS_WAREHOUSE_ID = "wh-shared"; // Re-initialize with the new env @@ -203,7 +263,7 @@ describe("ServiceContext", () => { const userCtx = ServiceContext.createUserContext("user-token", "user-1"); const serviceCtx = ServiceContext.get(); - expect(userCtx.warehouseId).toBe(serviceCtx.warehouseId); + expect(userCtx.warehouseId).toBe(getWarehouseId()); expect(userCtx.workspaceId).toBe(serviceCtx.workspaceId); }); @@ -288,6 +348,19 @@ describe("ServiceContext", () => { expect(state.serviceUserId).toBe("new-service-user"); }); + + test("clears app resources while preserving previously returned legacy snapshots", async () => { + process.env.DATABRICKS_WAREHOUSE_ID = "wh-before-reset"; + const previous = await ServiceContext.initialize({ warehouseId: true }); + const previousWarehouse = getWarehouseId(); + ServiceContext.reset(); + expect(getWarehouseId).toThrow(InitializationError); + process.env.DATABRICKS_WAREHOUSE_ID = "wh-after-reset"; + await ServiceContext.initialize({ warehouseId: true }); + expect(await getWarehouseId()).toBe("wh-after-reset"); + expect(previous.warehouseId).toBe(previousWarehouse); + expect(await previous.warehouseId).toBe("wh-before-reset"); + }); }); // ── getWorkspaceId() (private, tested via initialize) ───────── @@ -333,15 +406,15 @@ describe("ServiceContext", () => { }); }); - // ── getWarehouseId() (private, tested via initialize) ───────── + // Warehouse discovery remains covered through the compatible startup API. - describe("getWarehouseId()", () => { + describe("warehouse resource initialization", () => { test("should use DATABRICKS_WAREHOUSE_ID env var when set", async () => { process.env.DATABRICKS_WAREHOUSE_ID = "env-wh-abc"; - const state = await ServiceContext.initialize({ warehouseId: true }); + await ServiceContext.initialize({ warehouseId: true }); - expect(await state.warehouseId).toBe("env-wh-abc"); + expect(await getWarehouseId()).toBe("env-wh-abc"); }); test("should skip auto-resolve in dev mode when agentic mode is enabled", async () => { @@ -376,10 +449,10 @@ describe("ServiceContext", () => { return Promise.resolve({ "x-databricks-org-id": "ws-dev" }); }); - const state = await ServiceContext.initialize({ warehouseId: true }); + await ServiceContext.initialize({ warehouseId: true }); // Should pick RUNNING warehouse (highest priority) - expect(await state.warehouseId).toBe("wh-running"); + expect(await getWarehouseId()).toBe("wh-running"); }); test("should sort warehouses by state priority in dev mode", async () => { @@ -399,10 +472,10 @@ describe("ServiceContext", () => { return Promise.resolve({ "x-databricks-org-id": "ws-dev" }); }); - const state = await ServiceContext.initialize({ warehouseId: true }); + await ServiceContext.initialize({ warehouseId: true }); // STOPPED (priority 1) < STARTING (priority 2) < STOPPING (priority 3) - expect(await state.warehouseId).toBe("wh-stopped"); + expect(await getWarehouseId()).toBe("wh-stopped"); }); test("should throw in dev mode when no warehouses are available", async () => { diff --git a/packages/appkit/src/context/user-context.ts b/packages/appkit/src/context/user-context.ts index dddd9b4bd..b854f578e 100644 --- a/packages/appkit/src/context/user-context.ts +++ b/packages/appkit/src/context/user-context.ts @@ -1,8 +1,17 @@ +import { + type CallerContext, + type ExecutionContext, + isCallerContext, + snapshotCallerContext, +} from "./caller-context"; +import { warnContextDeprecation } from "./deprecation"; import type { ServiceContextState } from "./service-context"; +export type { ExecutionContext } from "./caller-context"; + /** - * User execution context extends the service context with user-specific data. - * Created on-demand when asUser(req) is called. + * @deprecated Use CallerContext and its principal field. Kept for callers + * that construct the legacy shape or read its flat identity fields. */ export interface UserContext { /** WorkspaceClient authenticated as the user */ @@ -15,7 +24,7 @@ export interface UserContext { userEmail?: string; /** Truncated SHA-256 hash of the user's OBO token, used to detect token rotation */ tokenFingerprint?: string; - /** Promise that resolves to the warehouse ID (inherited from service context, only present when a plugin requires `SQL_WAREHOUSE` resource) */ + /** @deprecated Use getWarehouseId() from @databricks/appkit. */ warehouseId?: Promise; /** Promise that resolves to the workspace ID (inherited from service context) */ workspaceId: Promise; @@ -24,13 +33,102 @@ export interface UserContext { } /** - * Execution context can be either service or user context. + * @deprecated Use snapshotCallerContext for identity-only snapshots. + * Retains the legacy fields for existing callers. */ -export type ExecutionContext = ServiceContextState | UserContext; +export function immutableCallerContext( + ctx: CallerContext, +): CallerContext & UserContext { + warnContextDeprecation("immutableCallerContext", "snapshotCallerContext"); + return legacyIdentityContext(ctx); +} + +function legacyIdentityContext( + ctx: CallerContext, +): CallerContext & UserContext { + const caller = snapshotCallerContext(ctx); + const { principal } = caller; + return Object.freeze({ + ...caller, + get userId() { + warnContextDeprecation( + "UserContext.userId", + "CallerContext.principal.userId", + ); + return principal.userId; + }, + get userName() { + warnContextDeprecation( + "UserContext.userName", + "CallerContext.principal.userName", + ); + return principal.userName; + }, + get userEmail() { + warnContextDeprecation( + "UserContext.userEmail", + "CallerContext.principal.userEmail", + ); + return principal.userEmail; + }, + get isUserContext(): true { + warnContextDeprecation( + "UserContext.isUserContext", + "CallerContext.principal.type", + ); + return true; + }, + }); +} + +/** Expose the old resource field only through the deprecated context APIs. */ +export function legacyUserContext( + ctx: CallerContext, + resolveWarehouseId: () => Promise | undefined, +): CallerContext & UserContext { + return Object.freeze( + Object.defineProperties( + { + get warehouseId() { + warnContextDeprecation( + "UserContext.warehouseId", + "getWarehouseId() from @databricks/appkit", + ); + return resolveWarehouseId(); + }, + }, + Object.getOwnPropertyDescriptors(legacyIdentityContext(ctx)), + ), + ) as CallerContext & UserContext; +} + +/** Normalize legacy inputs before opening a caller scope. */ +export function toCallerContext( + ctx: CallerContext | UserContext, +): CallerContext { + if ("principal" in ctx) return ctx; + return { + client: ctx.client, + principal: { + type: "user", + userId: ctx.userId, + userName: ctx.userName, + userEmail: ctx.userEmail, + }, + tokenFingerprint: ctx.tokenFingerprint, + workspaceId: ctx.workspaceId, + }; +} /** - * Check if an execution context is a user context. + * @deprecated Use isCallerContext. Active caller contexts retain the legacy + * identity accessors for callers narrowed by this guard. */ -export function isUserContext(ctx: ExecutionContext): ctx is UserContext { - return "isUserContext" in ctx && ctx.isUserContext === true; +export function isUserContext( + ctx: ExecutionContext | UserContext, +): ctx is UserContext & Partial { + warnContextDeprecation("isUserContext", "isCallerContext"); + return "principal" in ctx + ? isCallerContext(ctx) + : "isUserContext" in ctx && ctx.isUserContext === true; } diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 42227b992..aa29602b2 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -19,6 +19,7 @@ import { createLogger } from "../logging/logger"; import { isPlainObject } from "../plugin/plugin"; import { uiVariants } from "../plugins/ui-variants"; import { ResourceRegistry, ResourceType } from "../registry"; +import { getConfiguredWarehouseId } from "../resources/warehouse"; import type { TelemetryConfig } from "../telemetry"; import { TelemetryManager } from "../telemetry"; import type { WorkspaceClient } from "../workspace-client"; @@ -219,13 +220,17 @@ export class AppKit { // Collect manifest resources via registry const registry = new ResourceRegistry(); registry.collectResources(rawPlugins); + registry.validate(); // Derive ServiceContext needs from what manifests declared const needsWarehouse = registry .getRequired() .some((r) => r.type === ResourceType.SQL_WAREHOUSE); await ServiceContext.initialize( - { warehouseId: needsWarehouse }, + { + warehouseId: + getConfiguredWarehouseId(registry.getRequired()) ?? needsWarehouse, + }, config?.client, ); diff --git a/packages/appkit/src/core/tests/databricks.test.ts b/packages/appkit/src/core/tests/databricks.test.ts index f5cce386f..7cc1a8ae5 100644 --- a/packages/appkit/src/core/tests/databricks.test.ts +++ b/packages/appkit/src/core/tests/databricks.test.ts @@ -736,7 +736,7 @@ describe("AppKit", () => { initSpy.mockRestore(); }); - test("should call ServiceContext.initialize with warehouseId: true when a plugin requires sql_warehouse", async () => { + test("passes the warehouse ID resolved from the plugin's declared environment variable", async () => { const PluginWithRequiredResource = class extends CoreTestPlugin { static manifest: PluginManifest = { name: "withResource", @@ -750,15 +750,15 @@ describe("AppKit", () => { resourceKey: "warehouse", description: "Warehouse", permission: "CAN_USE", - fields: { id: { env: "DATABRICKS_WAREHOUSE_ID" } }, + fields: { id: { env: "APPKIT_TEST_WAREHOUSE_ID" } }, }, ], optional: [], }, }; }; - const prevWh = process.env.DATABRICKS_WAREHOUSE_ID; - process.env.DATABRICKS_WAREHOUSE_ID = "wh-123"; + vi.stubEnv("APPKIT_TEST_WAREHOUSE_ID", "wh-plugin"); + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", "wh-unrelated"); try { const contextModule = await import("../../context/service-context"); const initSpy = vi.spyOn(contextModule.ServiceContext, "initialize"); @@ -771,11 +771,13 @@ describe("AppKit", () => { }, ], }); - expect(initSpy).toHaveBeenCalledWith({ warehouseId: true }, undefined); + expect(initSpy).toHaveBeenCalledWith( + { warehouseId: "wh-plugin" }, + undefined, + ); initSpy.mockRestore(); } finally { - if (prevWh !== undefined) process.env.DATABRICKS_WAREHOUSE_ID = prevWh; - else delete process.env.DATABRICKS_WAREHOUSE_ID; + vi.unstubAllEnvs(); } }); }); diff --git a/packages/appkit/src/index.ts b/packages/appkit/src/index.ts index 854db3b5f..57848eeb0 100644 --- a/packages/appkit/src/index.ts +++ b/packages/appkit/src/index.ts @@ -36,8 +36,16 @@ export { getWorkspaceClient, RequestedClaimsPermissionSet, } from "./connectors/lakebase"; -export { getExecutionContext } from "./context"; +export { + type CallerContext, + type CallerPrincipal, + type ExecutionContext, + getCurrentActorId, + getCurrentPrincipalKey, + getExecutionContext, +} from "./context"; export { createApp } from "./core"; +export { getWarehouseId } from "./resources"; export type { DatabaseRegistry } from "./database/contract"; // Errors export { diff --git a/packages/appkit/src/plugin/plugin.ts b/packages/appkit/src/plugin/plugin.ts index a636aa61c..a9ea78cf9 100644 --- a/packages/appkit/src/plugin/plugin.ts +++ b/packages/appkit/src/plugin/plugin.ts @@ -16,7 +16,11 @@ import { camelToKebab } from "shared"; import { AppManager } from "../app"; import { CacheManager } from "../cache"; -import { getCurrentUserId, runInUserContext, ServiceContext } from "../context"; +import { + getCurrentUserId, + runInCallerContext, + ServiceContext, +} from "../context"; import type { PluginContext } from "../core/plugin-context"; import { AppKitError, AuthenticationError } from "../errors"; import { createLogger } from "../logging/logger"; @@ -406,7 +410,7 @@ export abstract class Plugin< * * Returns the `x-forwarded-user` header when present. In development mode * (`NODE_ENV=development`) falls back to the current context user ID so - * that callers outside an active `runInUserContext` scope still get a + * that callers outside an active `runInCallerContext` scope still get a * consistent value. * * @throws AuthenticationError in production when no user header is present. @@ -458,7 +462,7 @@ export abstract class Plugin< const effectiveUserId = userId || "dev-user"; - const userContext = ServiceContext.createUserContext( + const userContext = ServiceContext.createCallerContext( token, effectiveUserId, undefined, @@ -468,7 +472,7 @@ export abstract class Plugin< return this._createAsUserProxy( (fn) => (...args) => - runInUserContext(userContext, () => fn(...args)), + runInCallerContext(userContext, () => fn(...args)), ); } @@ -477,7 +481,7 @@ export abstract class Plugin< * 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)) + * - real OBO: fn => (...args) => runInCallerContext(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 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 9f8cf3d10..3f38539a5 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 @@ -14,11 +14,10 @@ vi.mock("../../../context", () => ({ getCurrentUserId: vi.fn(() => "test-user"), // OBO plumbing so asUser() runs its non-dev path. getCurrentUserId stays // constant, so per-user scoping is driven by executorKey in the cacheKey. - runInUserContext: (_ctx: unknown, fn: () => T): T => fn(), + runInCallerContext: (_ctx: unknown, fn: () => T): T => fn(), ServiceContext: { - createUserContext: (_token: string, userId: string) => ({ - userId, - isUserContext: true, + createCallerContext: (_token: string, userId: string) => ({ + principal: { type: "user", userId }, }), }, })); diff --git a/packages/appkit/src/plugins/analytics/analytics.ts b/packages/appkit/src/plugins/analytics/analytics.ts index dc3543be4..987ac51d3 100644 --- a/packages/appkit/src/plugins/analytics/analytics.ts +++ b/packages/appkit/src/plugins/analytics/analytics.ts @@ -16,7 +16,7 @@ import { DEFAULT_WAREHOUSE_STARTUP_TIMEOUT_MS, type WarehouseStatusUpdate, } from "../../connectors/sql-warehouse/client"; -import { getWarehouseId, getWorkspaceClient } from "../../context"; +import { getWorkspaceClient } from "../../context"; import { buildToolkitEntries } from "../../core/agent/build-toolkit"; import { defineTool, @@ -28,6 +28,7 @@ import { AppKitError, ExecutionError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; +import { getWarehouseId } from "../../resources"; import type { WorkspaceClient } from "../../workspace-client"; import { queryDefaults } from "./defaults"; import manifest from "./manifest.json"; diff --git a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts index e3f71480c..bb42cb830 100644 --- a/packages/appkit/src/plugins/analytics/tests/analytics.test.ts +++ b/packages/appkit/src/plugins/analytics/tests/analytics.test.ts @@ -130,7 +130,9 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); // Verify service workspace client is used - expect(capturedWorkspaceClient).toBeDefined(); + expect(capturedWorkspaceClient).toBe( + serviceContextMock.serviceContext.client, + ); // Verify executeStatement is called with correct statement expect(executeMock).toHaveBeenCalledWith( @@ -197,8 +199,13 @@ describe("Analytics Plugin", () => { await handler(mockReq, mockRes); - // Verify a workspace client is used - expect(capturedWorkspaceClient).toBeDefined(); + const callerContext = + serviceContextMock.createUserContextSpy.mock.results[0].value; + expect(callerContext).not.toHaveProperty("warehouseId"); + expect(capturedWorkspaceClient).toBe(callerContext.client); + expect(capturedWorkspaceClient).not.toBe( + serviceContextMock.serviceContext.client, + ); // Verify the query is executed with correct statement expect(executeMock).toHaveBeenCalledWith( diff --git a/packages/appkit/src/plugins/files/plugin.ts b/packages/appkit/src/plugins/files/plugin.ts index fa175870c..cc3dc4ee8 100644 --- a/packages/appkit/src/plugins/files/plugin.ts +++ b/packages/appkit/src/plugins/files/plugin.ts @@ -22,11 +22,11 @@ import { getCurrentUserId, getExecutionContext, getWorkspaceClient, - runInUserContext, + runInCallerContext, ServiceContext, - type UserContext, + type CallerContext, + isCallerContext, } from "../../context"; -import { isUserContext } from "../../context/user-context"; import { buildToolkitEntries } from "../../core/agent/build-toolkit"; import { defineTool, @@ -163,7 +163,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { * return a policy user explicitly marked `isServicePrincipal: true`, so * even in dev a `usersOnly`-style policy that gates on * `!user.isServicePrincipal` cannot be tricked. The matching SDK execution - * path also falls through to the SP client (no `runInUserContext` wrap), + * path also falls through to the SP client (no `runInCallerContext` wrap), * so the policy user and the SDK identity stay aligned. */ private _extractUser(req: express.Request): FilePolicyUser { @@ -265,7 +265,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { * NOTE: This method only selects which identity the *policy* sees. The * matching SDK execution identity is selected separately by * `_resolveAuthForRequest` and applied via `_runWithAuth` / - * `runInUserContext` in each handler. The two selections are designed to + * `runInCallerContext` in each handler. The two selections are designed to * converge on the same identity per the policy-user matrix in the docs — * see `docs/docs/plugins/files.md#policy-user-matrix`. */ @@ -1233,7 +1233,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { const settings = this._writeSettings(mode); // The connector's `upload` resolves `getWorkspaceClient()` and // `client.config.authenticate(headers)` synchronously inside this - // callback. When `_runWithAuth` wraps us in `runInUserContext`, that + // callback. When `_runWithAuth` wraps us in `runInCallerContext`, that // chain produces user-token Authorization headers on the outgoing // `fetch PUT`. The OBO upload-headers test pins this contract. const result = await this.trackWrite(() => @@ -1386,7 +1386,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { } /** - * Build a `UserContext` from request headers when both + * Build a `CallerContext` from request headers when both * `x-forwarded-access-token` and `x-forwarded-user` are present, otherwise * return `null`. Used by OBO route handlers to wrap SDK calls in the * end-user's identity. A `null` result means "fall back to the service @@ -1394,17 +1394,17 @@ export class FilesPlugin extends Plugin implements ToolProvider { * already have responded 401 before we get here, so `null` is reachable * only on the dev-fallback path. */ - private _buildUserContextOrNull(req: express.Request): UserContext | null { + private _buildUserContextOrNull(req: express.Request): CallerContext | null { const token = req.header("x-forwarded-access-token")?.trim(); const userId = req.header("x-forwarded-user")?.trim(); if (!token || !userId) return null; - return ServiceContext.createUserContext(token, userId); + return ServiceContext.createCallerContext(token, userId); } /** * Build the telemetry attribute hash for the `files.auth_mode` span * attribute. The value reflects what operationally happened — i.e. - * whether `runInUserContext` actually wrapped the SDK call: + * whether `runInCallerContext` actually wrapped the SDK call: * - HTTP route on OBO volume + valid token → `"on-behalf-of-user"`. * - HTTP route on OBO volume + dev-fallback (no token) → * `"service-principal"` (the route falls through to the SP client). @@ -1421,12 +1421,12 @@ export class FilesPlugin extends Plugin implements ToolProvider { /** * One-shot resolver for HTTP route handlers. Builds the request's - * `UserContext` AT MOST ONCE (when the volume is OBO and the headers are + * `CallerContext` AT MOST ONCE (when the volume is OBO and the headers are * present) and returns both the operationally-effective auth mode and the - * pre-built `UserContext`. + * pre-built `CallerContext`. * * Handlers thread the `userCtx` into `_runWithAuth(userCtx, fn)` to avoid - * a second `ServiceContext.createUserContext()` allocation — that call + * a second `ServiceContext.createCallerContext()` allocation. That call * builds a fresh `WorkspaceClient` per invocation, so doing it twice per * request was pure throwaway overhead. */ @@ -1435,7 +1435,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { volumeKey: string, ): { mode: "service-principal" | "on-behalf-of-user"; - userCtx: UserContext | null; + userCtx: CallerContext | null; } { if (this._resolveAuth(volumeKey) !== "on-behalf-of-user") { return { mode: "service-principal", userCtx: null }; @@ -1452,7 +1452,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { * `WorkspaceClient` and `getCurrentUserId()` 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 `UserContext`: wraps `fn` in `runInUserContext(userCtx)`, + * - `userCtx` is a `CallerContext`: wraps `fn` in `runInCallerContext(userCtx)`, * so SDK calls execute as the end user and `getCurrentUserId()` (and * therefore cache keys) resolve to the user's ID. * @@ -1461,11 +1461,11 @@ export class FilesPlugin extends Plugin implements ToolProvider { * NOT take a `req` so it cannot accidentally re-build the context. */ private async _runWithAuth( - userCtx: UserContext | null, + userCtx: CallerContext | null, fn: () => Promise, ): Promise { if (userCtx) { - return runInUserContext(userCtx, fn); + return runInCallerContext(userCtx, fn); } return fn(); } @@ -1529,7 +1529,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { /** * Wrap each `VolumeAPI` method so its execution runs inside - * `runInUserContext(userCtx, ...)`. Used by `VolumeHandle.asUser(req)` to + * `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()` @@ -1541,7 +1541,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { */ private _wrapVolumeAPIInUserContext( api: VolumeAPI, - userCtx: UserContext, + userCtx: CallerContext, ): VolumeAPI { const wrap = ( @@ -1550,7 +1550,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { ): ((...args: Args) => Promise) => (...args: Args) => this._withAuthModeAttributes(operation, "on-behalf-of-user", () => - runInUserContext(userCtx, () => fn(...args)), + runInCallerContext(userCtx, () => fn(...args)), ); return { @@ -1639,8 +1639,8 @@ export class FilesPlugin extends Plugin implements ToolProvider { private _defineVolumeTools(volumeKey: string): ToolRegistry { const buildUser = (): FilePolicyUser => { const ctx = getExecutionContext(); - return isUserContext(ctx) - ? { id: ctx.userId } + return isCallerContext(ctx) + ? { id: ctx.principal.userId } : { id: ctx.serviceUserId, isServicePrincipal: true }; }; const api = () => this.createVolumeAPI(volumeKey, buildUser()); @@ -1781,7 +1781,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { * through the HTTP routes run as the end user; for programmatic calls * outside a route, use `asUser(req)` to opt into per-user execution. * `asUser(req)` is a hard override at the SDK level: it forces every - * subsequent call to execute as the end user inside `runInUserContext`, + * subsequent call to execute as the end user inside `runInCallerContext`, * regardless of the volume's `auth` setting. Policies control per-user * access in either mode. * @@ -1824,7 +1824,7 @@ export class FilesPlugin extends Plugin implements ToolProvider { const user = this._extractUser(req); const api = this.createVolumeAPI(volumeKey, user); // Force OBO at the SDK level regardless of the volume's `auth` - // setting: each method runs inside `runInUserContext` so + // setting: each method runs inside `runInCallerContext` so // `getWorkspaceClient()` returns the user-token client. When no // user token is available (only reachable in dev mode after the // strict `_extractUser` falls back to the SP identity), we skip diff --git a/packages/appkit/src/plugins/files/tests/plugin.test.ts b/packages/appkit/src/plugins/files/tests/plugin.test.ts index a6736d307..ad0cbc4aa 100644 --- a/packages/appkit/src/plugins/files/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/files/tests/plugin.test.ts @@ -2738,10 +2738,8 @@ describe("FilesPlugin", () => { serviceContextMock.createUserContextSpy.mockImplementation( (_token: string, userId: string) => ({ client: userClient as any, - userId, - warehouseId: serviceContextMock.serviceContext.warehouseId, + principal: { type: "user", userId }, workspaceId: serviceContextMock.serviceContext.workspaceId, - isUserContext: true, }), ); @@ -2914,10 +2912,8 @@ describe("FilesPlugin", () => { serviceContextMock.createUserContextSpy.mockImplementation( (_token: string, userId: string) => ({ client: userClient as any, - userId, - warehouseId: serviceContextMock.serviceContext.warehouseId, + principal: { type: "user", userId }, workspaceId: serviceContextMock.serviceContext.workspaceId, - isUserContext: true, }), ); @@ -3348,10 +3344,8 @@ describe("FilesPlugin", () => { serviceContextMock.createUserContextSpy.mockImplementation( (_token: string, userId: string) => ({ client: userClient as any, - userId, - warehouseId: serviceContextMock.serviceContext.warehouseId, + principal: { type: "user", userId }, workspaceId: serviceContextMock.serviceContext.workspaceId, - isUserContext: true, }), ); diff --git a/packages/appkit/src/plugins/lakebase/lakebase.ts b/packages/appkit/src/plugins/lakebase/lakebase.ts index 8518b3ab2..11e7889d6 100644 --- a/packages/appkit/src/plugins/lakebase/lakebase.ts +++ b/packages/appkit/src/plugins/lakebase/lakebase.ts @@ -13,7 +13,7 @@ import { RoutingPool, } from "../../connectors/lakebase"; import { getClientOptions } from "../../context/client-options"; -import { getUserContext } from "../../context/execution-context"; +import { getCallerContext } from "../../context/execution-context"; import { buildToolkitEntries } from "../../core/agent/build-toolkit"; import { defineTool, @@ -104,7 +104,7 @@ export class LakebasePlugin extends Plugin implements ToolProvider { this.pool = new RoutingPool(spPool, (ctx) => { if (!oboManager) throw new Error("OBO pool manager not initialized"); // Lakebase OAuth roles use email as the postgres role when available - const userKey = ctx.userEmail ?? ctx.userId; + const userKey = ctx.principal.userEmail ?? ctx.principal.userId; const isNew = !oboManager.hasPool(userKey); const pool = oboManager.getPool( userKey, @@ -299,9 +299,9 @@ export class LakebasePlugin extends Plugin implements ToolProvider { * Inside `asUser(req)`, returns user-scoped config; otherwise SP config. */ private activePoolConfig() { - const ctx = getUserContext(); + const ctx = getCallerContext(); if (ctx) { - const user = ctx.userEmail ?? ctx.userId; + const user = ctx.principal.userEmail ?? ctx.principal.userId; return { ...this.config.pool, workspaceClient: ctx.client.toLegacyWorkspaceClient(), diff --git a/packages/appkit/src/resources/app-resources.ts b/packages/appkit/src/resources/app-resources.ts new file mode 100644 index 000000000..9a329ab2a --- /dev/null +++ b/packages/appkit/src/resources/app-resources.ts @@ -0,0 +1,40 @@ +import { warnContextDeprecation } from "../context/deprecation"; +import { InitializationError } from "../errors"; +import type { WorkspaceClient } from "../workspace-client"; +import { WarehouseResource, type WarehouseBinding } from "./warehouse"; + +/** @deprecated Use WarehouseBinding. Retained for backward compatibility. */ +export type AppResourceBindings = WarehouseBinding; + +/** @deprecated Use WarehouseResource. Retained for backward compatibility. */ +export class AppResources { + static resolve( + client: WorkspaceClient, + options?: { warehouseId?: boolean }, + ): Promise { + warnContextDeprecation("AppResources.resolve", "WarehouseResource.resolve"); + return WarehouseResource.resolve(client, options?.warehouseId); + } + + static bind(bindings: AppResourceBindings): AppResourceBindings { + warnContextDeprecation("AppResources.bind", "WarehouseResource.bind"); + return WarehouseResource.bind(bindings); + } + + static get(): AppResourceBindings { + warnContextDeprecation("AppResources.get", "WarehouseResource.get"); + const binding = WarehouseResource.get(); + if (!binding) { + throw InitializationError.notInitialized( + "AppResources", + "Call createApp() first", + ); + } + return binding; + } + + static reset(): void { + warnContextDeprecation("AppResources.reset", "WarehouseResource.reset"); + WarehouseResource.reset(); + } +} diff --git a/packages/appkit/src/resources/index.ts b/packages/appkit/src/resources/index.ts new file mode 100644 index 000000000..dd21545f1 --- /dev/null +++ b/packages/appkit/src/resources/index.ts @@ -0,0 +1 @@ +export { getWarehouseId } from "./warehouse"; diff --git a/packages/appkit/src/resources/tests/warehouse.test.ts b/packages/appkit/src/resources/tests/warehouse.test.ts new file mode 100644 index 000000000..7bc134519 --- /dev/null +++ b/packages/appkit/src/resources/tests/warehouse.test.ts @@ -0,0 +1,186 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { ServiceContext } from "../../context/service-context"; +import { ConfigurationError, InitializationError } from "../../errors"; +import { getWarehouseId } from "../../index"; +import { ResourceRegistry, ResourceType } from "../../registry"; +import { createMockWorkspaceClient } from "../../testing/mock-workspace-client"; +import { AppResources, type AppResourceBindings } from "../app-resources"; +import { getConfiguredWarehouseId, WarehouseResource } from "../warehouse"; + +function warehouseRegistry(bindings: { env: string; required?: boolean }[]) { + const registry = new ResourceRegistry(); + for (const { env, required = true } of bindings) { + registry.register("custom", { + type: ResourceType.SQL_WAREHOUSE, + resourceKey: env, + alias: "Query warehouse", + description: "Test warehouse", + permission: "CAN_USE", + required, + fields: { id: { env } }, + }); + } + registry.validate(); + return registry; +} + +describe("warehouse resource bindings", () => { + beforeEach(() => { + WarehouseResource.reset(); + }); + + afterEach(() => { + WarehouseResource.reset(); + vi.restoreAllMocks(); + vi.unstubAllEnvs(); + }); + + test("the public accessor requires initialized resources", () => { + expect(getWarehouseId).toThrow(InitializationError); + }); + + test("reads the resource binding without consulting service identity", () => { + const warehouseId = Promise.resolve("resource-warehouse"); + WarehouseResource.bind({ warehouseId }); + const getService = vi + .spyOn(ServiceContext, "get") + .mockImplementation(() => { + throw new Error("Service identity is not available"); + }); + expect(getWarehouseId()).toBe(warehouseId); + expect(getService).not.toHaveBeenCalled(); + }); + + test("freezes a resource snapshot without freezing its input", () => { + const warehouseId = Promise.resolve("original"); + const input = { warehouseId }; + const bindings = WarehouseResource.bind(input); + input.warehouseId = Promise.resolve("changed"); + expect(Object.isFrozen(bindings)).toBe(true); + expect(Reflect.set(bindings, "warehouseId", input.warehouseId)).toBe(false); + expect(getWarehouseId()).toBe(warehouseId); + }); + + test("resolves the environment binding without publishing it", async () => { + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", "configured-warehouse"); + const client = createMockWorkspaceClient(); + const bindings = await WarehouseResource.resolve(client, true); + expect(await bindings.warehouseId).toBe("configured-warehouse"); + expect(client.apiClient.request).not.toHaveBeenCalled(); + expect(getWarehouseId).toThrow(InitializationError); + WarehouseResource.bind(bindings); + expect(getWarehouseId()).toBe(bindings.warehouseId); + }); + + test("does not resolve a warehouse when no plugin requires it", async () => { + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", "unused-warehouse"); + const client = createMockWorkspaceClient(); + const bindings = await WarehouseResource.resolve(client); + expect(bindings.warehouseId).toBeUndefined(); + expect(client.apiClient.request).not.toHaveBeenCalled(); + WarehouseResource.bind(bindings); + expect(getWarehouseId).toThrow(ConfigurationError); + expect(getWarehouseId).toThrow("No plugin requires a SQL Warehouse"); + }); + + test("uses the manifest binding instead of an unrelated default environment variable", async () => { + vi.stubEnv("APPKIT_TEST_WAREHOUSE", "plugin-warehouse"); + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", "unrelated-warehouse"); + const registry = warehouseRegistry([{ env: "APPKIT_TEST_WAREHOUSE" }]); + const client = createMockWorkspaceClient(); + const binding = await WarehouseResource.resolve( + client, + getConfiguredWarehouseId(registry.getAll()), + ); + WarehouseResource.bind(binding); + expect(await getWarehouseId()).toBe("plugin-warehouse"); + expect(client.apiClient.request).not.toHaveBeenCalled(); + }); + + test("shares one ID across declarations with different environment variables", () => { + vi.stubEnv("APPKIT_TEST_WAREHOUSE", "shared-warehouse"); + vi.stubEnv("APPKIT_TEST_OTHER_WAREHOUSE", "shared-warehouse"); + const registry = warehouseRegistry([ + { env: "APPKIT_TEST_WAREHOUSE" }, + { env: "APPKIT_TEST_OTHER_WAREHOUSE" }, + ]); + expect(getConfiguredWarehouseId(registry.getAll())).toBe( + "shared-warehouse", + ); + }); + + test("requires an explicit app default when manifest bindings differ", () => { + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", ""); + vi.stubEnv("APPKIT_TEST_WAREHOUSE", "first-warehouse"); + vi.stubEnv("APPKIT_TEST_OTHER_WAREHOUSE", "second-warehouse"); + const registry = warehouseRegistry([ + { env: "APPKIT_TEST_WAREHOUSE" }, + { env: "APPKIT_TEST_OTHER_WAREHOUSE" }, + ]); + expect(() => getConfiguredWarehouseId(registry.getAll())).toThrow( + "Multiple SQL warehouses are configured. Set DATABRICKS_WAREHOUSE_ID", + ); + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", "second-warehouse"); + expect(getConfiguredWarehouseId(registry.getAll())).toBe( + "second-warehouse", + ); + }); + + test("leaves missing bindings to discovery and does not activate optional warehouses", () => { + vi.stubEnv("APPKIT_TEST_WAREHOUSE", ""); + vi.stubEnv("APPKIT_TEST_OTHER_WAREHOUSE", "optional-warehouse"); + const registry = warehouseRegistry([ + { env: "APPKIT_TEST_WAREHOUSE" }, + { env: "APPKIT_TEST_OTHER_WAREHOUSE", required: false }, + ]); + expect(getConfiguredWarehouseId(registry.getAll())).toBeUndefined(); + }); + + test("deprecated resource names delegate to the warehouse module and warn once", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", "compat-warehouse"); + const client = createMockWorkspaceClient(); + for (let i = 0; i < 2; i++) { + expect(() => AppResources.get()).toThrow(InitializationError); + const binding: AppResourceBindings = await AppResources.resolve(client, { + warehouseId: true, + }); + expect(AppResources.bind(binding)).toBe(WarehouseResource.get()); + expect(AppResources.get()).toBe(WarehouseResource.get()); + expect(await getWarehouseId()).toBe("compat-warehouse"); + AppResources.reset(); + expect(getWarehouseId).toThrow(InitializationError); + } + for (const method of ["resolve", "bind", "get", "reset"]) { + expect( + warn.mock.calls.filter((args) => + args.includes( + `AppResources.${method} is deprecated. Use WarehouseResource.${method} instead.`, + ), + ), + ).toHaveLength(1); + } + }); + // Databricks list APIs omit empty arrays (`warehouses` is `omitempty`), so + // dev discovery must treat a missing key like an empty list. + test.each([ + { name: "omitted warehouses key", response: {} }, + { name: "empty warehouses list", response: { warehouses: [] } }, + ])( + "dev discovery reports a missing warehouse for $name", + async ({ response }) => { + vi.stubEnv("NODE_ENV", "development"); + vi.stubEnv("DATABRICKS_WAREHOUSE_ID", ""); + vi.stubEnv("DATABRICKS_APPS_AGENTIC_MODE", ""); + const client = createMockWorkspaceClient(); + vi.mocked(client.apiClient.request).mockResolvedValue(response); + const error = await WarehouseResource.resolve(client, true).catch( + (e: unknown) => e, + ); + expect(error).toBeInstanceOf(ConfigurationError); + expect(error).not.toBeInstanceOf(TypeError); + expect(client.apiClient.request).toHaveBeenCalledTimes(1); + }, + ); +}); diff --git a/packages/appkit/src/resources/warehouse.ts b/packages/appkit/src/resources/warehouse.ts new file mode 100644 index 000000000..461f24699 --- /dev/null +++ b/packages/appkit/src/resources/warehouse.ts @@ -0,0 +1,176 @@ +import { AsyncLocalStorage } from "node:async_hooks"; + +import { ConfigurationError, InitializationError } from "../errors"; +import type { ResourceEntry } from "../registry/types"; +import type { sql, WorkspaceClient } from "../workspace-client"; + +/** The SQL warehouse selected for the app lifecycle, separate from identity. */ +export interface WarehouseBinding { + readonly warehouseId?: Promise; +} + +let appBinding: WarehouseBinding | undefined; + +/** Owns SQL warehouse discovery and the app's immutable warehouse binding. */ +export class WarehouseResource { + /** Resolve candidates without publishing resources from a failed app startup. */ + static async resolve( + client: WorkspaceClient, + selection: boolean | string = false, + ): Promise { + const warehouseId = + typeof selection === "string" + ? selection + : selection + ? await discoverWarehouseId(client) + : undefined; + return Object.freeze({ + warehouseId: + warehouseId === undefined ? undefined : Promise.resolve(warehouseId), + }); + } + + /** Publish an immutable snapshot of resolved app resources. */ + static bind(bindings: WarehouseBinding): WarehouseBinding { + appBinding = Object.freeze({ ...bindings }); + return appBinding; + } + + static get(): WarehouseBinding | undefined { + return appBinding; + } + + /** Reset app bindings alongside the service context in tests. */ + static reset(): void { + appBinding = undefined; + } +} + +/** Select the app warehouse from resource values resolved by the registry. */ +export function getConfiguredWarehouseId( + resources: readonly ResourceEntry[], +): string | undefined { + const ids = new Set( + resources + .filter( + (resource) => resource.type === "sql_warehouse" && resource.required, + ) + .flatMap((resource) => (resource.values?.id ? [resource.values.id] : [])), + ); + if (ids.size <= 1) return ids.values().next().value; + + // The shared accessor needs an explicit default when bindings differ. + if (process.env.DATABRICKS_WAREHOUSE_ID) { + return process.env.DATABRICKS_WAREHOUSE_ID; + } + throw new ConfigurationError( + "Multiple SQL warehouses are configured. Set DATABRICKS_WAREHOUSE_ID to select the app's default warehouse.", + ); +} + +async function discoverWarehouseId(client: WorkspaceClient): Promise { + if (process.env.DATABRICKS_WAREHOUSE_ID) { + return process.env.DATABRICKS_WAREHOUSE_ID; + } + + const agenticMode = + process.env.DATABRICKS_APPS_AGENTIC_MODE === "true" || + process.env.DATABRICKS_APPS_AGENTIC_MODE === "1"; + + if (process.env.NODE_ENV === "development" && !agenticMode) { + const response = (await client.apiClient.request({ + path: "/api/2.0/sql/warehouses", + method: "GET", + headers: new Headers(), + raw: false, + query: { skip_cannot_use: "true" }, + })) as { warehouses: sql.EndpointInfo[] }; + + const priorities: Record = { + RUNNING: 0, + STOPPED: 1, + STARTING: 2, + STOPPING: 3, + DELETED: 99, + DELETING: 99, + }; + + const warehouses = (response.warehouses || []).sort((a, b) => { + return ( + priorities[a.state as sql.State] - priorities[b.state as sql.State] + ); + }); + + if (warehouses.length === 0) { + throw ConfigurationError.resourceNotFound( + "Warehouse ID", + "Please configure the DATABRICKS_WAREHOUSE_ID environment variable", + ); + } + + const firstWarehouse = warehouses[0]; + if ( + firstWarehouse.state === "DELETED" || + firstWarehouse.state === "DELETING" || + !firstWarehouse.id + ) { + throw ConfigurationError.resourceNotFound( + "Warehouse ID", + "Please configure the DATABRICKS_WAREHOUSE_ID environment variable", + ); + } + + return firstWarehouse.id; + } + + throw ConfigurationError.resourceNotFound( + "Warehouse ID", + "Please configure the DATABRICKS_WAREHOUSE_ID environment variable", + ); +} + +// Only deprecated user-context overrides use this scope. New callers share app bindings. +const legacyResourceStorage = new AsyncLocalStorage< + WarehouseBinding | undefined +>(); + +/** @internal Keep legacy resource overrides out of execution identity. */ +export function runWithResourceBindings( + bindings: WarehouseBinding | undefined, + fn: () => T, +): T { + return legacyResourceStorage.run(bindings, fn); +} + +/** + * Get the configured SQL warehouse ID after app initialization. + * The warehouse is an app resource; SP and caller executions use the same binding. + * Deprecated user-context scopes retain support for explicit warehouse overrides. + * + * @throws ConfigurationError if no SQL warehouse was required at startup. + * @throws InitializationError if the app resources are not initialized. + */ +export function getWarehouseId(): Promise { + const warehouseId = captureWarehouseId()(); + if (!warehouseId) { + throw ConfigurationError.resourceNotFound( + "Warehouse ID", + "No plugin requires a SQL Warehouse. Add a sql_warehouse resource to your plugin manifest, or set DATABRICKS_WAREHOUSE_ID", + ); + } + return warehouseId; +} + +/** Capture compatibility access now so it cannot drift into a later scope. */ +export function captureWarehouseId(): () => Promise | undefined { + const binding = legacyResourceStorage.getStore() ?? WarehouseResource.get(); + return () => { + if (!binding) { + throw InitializationError.notInitialized( + "WarehouseResource", + "Call createApp() first", + ); + } + return binding.warehouseId; + }; +} diff --git a/packages/appkit/src/testing/create-test-app.ts b/packages/appkit/src/testing/create-test-app.ts index 159646980..fb4decdee 100644 --- a/packages/appkit/src/testing/create-test-app.ts +++ b/packages/appkit/src/testing/create-test-app.ts @@ -139,14 +139,14 @@ export interface TestApp { } /** - * Point `ServiceContext.createUserContext` at the harness's mock so an `obo` + * Point `ServiceContext.createCallerContext` at the harness's mock so an `obo` * request does not construct a real SDK client from `DATABRICKS_HOST`. * * Mirrors the `createUserContextSpy` in `fixtures.ts`; returns its restore. */ function stubUserContext(client: WorkspaceClient): () => void { const spy = vi - .spyOn(ServiceContext, "createUserContext") + .spyOn(ServiceContext, "createCallerContext") .mockImplementation((token, userId, userName, userEmail) => fakeUserContext(client, ServiceContext.get())( token, @@ -284,7 +284,7 @@ export async function createTestApp( suppliedClient ?? createMockWorkspaceClient({ responses, strict }); // createApp({ client }) installs only the service-principal client. An `obo` - // request reaches ServiceContext.createUserContext, which builds a *real* + // request reaches ServiceContext.createCallerContext, which builds a *real* // client from process.env.DATABRICKS_HOST — so the user-scoped path is faked // here too, or "no network" is false the moment a handler calls asUser. restoreUserContext = stubUserContext(client); diff --git a/packages/appkit/src/testing/fixtures.ts b/packages/appkit/src/testing/fixtures.ts index 97d93f8b7..5570c90f7 100644 --- a/packages/appkit/src/testing/fixtures.ts +++ b/packages/appkit/src/testing/fixtures.ts @@ -5,9 +5,14 @@ import type { IAppRouter } from "shared"; import { afterEach, beforeEach, vi } from "vitest"; import { CacheManager } from "../cache"; +import { snapshotCallerContext } from "../context/caller-context"; import type { ServiceContextState } from "../context/service-context"; import { ServiceContext } from "../context/service-context"; import { AuthenticationError } from "../errors"; +import { + WarehouseResource, + type WarehouseBinding, +} from "../resources/warehouse"; import type { InstrumentConfig, ITelemetry } from "../telemetry/types"; import { ApiError } from "../workspace-client"; import { createMockWorkspaceClient } from "./mock-workspace-client"; @@ -127,7 +132,7 @@ export type OboOption = }; /** - * The one fake of `ServiceContext.createUserContext` this kit uses, shared by + * The one fake of `ServiceContext.createCallerContext` this kit uses, shared by * `mockServiceContext` and `createTestApp`. * * Shared rather than duplicated because the two used to disagree, and neither @@ -138,10 +143,7 @@ export type OboOption = * * @internal */ -export function fakeUserContext( - client: Any, - ids: { warehouseId?: Any; workspaceId: Any }, -) { +export function fakeUserContext(client: Any, ids: { workspaceId: Any }) { return ( token: string, userId: string, @@ -151,21 +153,17 @@ export function fakeUserContext( // Same rejection as production, so a path that forgets to forward the token // fails here instead of only in a deployed app. if (!token) throw AuthenticationError.missingToken("user token"); - return { + return snapshotCallerContext({ client, - userId, - userName, - userEmail, + principal: { type: "user", userId, userName, userEmail }, // Derived from the token exactly as production does. Keyed on the user it // would be constant across tokens, and rotation compares this value. tokenFingerprint: createHash("sha256") .update(token) .digest("hex") .slice(0, 16), - warehouseId: ids.warehouseId, workspaceId: ids.workspaceId, - isUserContext: true, - }; + }); }; } @@ -491,27 +489,37 @@ export interface TestContextOptions { * installs the state as spies — that installer is the public entry point. */ function buildServiceContextState( - options: TestContextOptions = {}, + options: TestContextOptions, + resources: WarehouseBinding, ): ServiceContextState { return { client: (options.serviceDatabricksClient || createMockWorkspaceClient()) as Any, serviceUserId: options.serviceUserId || "test-service-user", - warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), + // Preserve the deprecated field for external test callers during migration. + warehouseId: resources.warehouseId, workspaceId: Promise.resolve(options.workspaceId || "test-workspace-id"), }; } /** * Mocks the `ServiceContext` singleton for testing — spies `get`, - * `initialize`, `isInitialized`, and `createUserContext` so code that resolves + * `initialize`, `isInitialized`, and `createCallerContext` so code that resolves * the service principal or an on-behalf-of user context gets test doubles. + * Also supplies the app-level warehouse binding through WarehouseResource. * Call in `beforeEach`; call the returned `restore()` in `afterEach`. * * @returns The mock context plus the spies and a `restore()` helper. */ export function mockServiceContext(options: TestContextOptions = {}) { - const serviceContext = buildServiceContextState(options); + const resources = Object.freeze({ + warehouseId: Promise.resolve(options.warehouseId || "test-warehouse-id"), + }); + const serviceContext = buildServiceContextState(options, resources); + + const resourcesSpy = vi + .spyOn(WarehouseResource, "get") + .mockReturnValue(resources); const getSpy = vi .spyOn(ServiceContext, "get") @@ -525,8 +533,10 @@ export function mockServiceContext(options: TestContextOptions = {}) { .spyOn(ServiceContext, "isInitialized") .mockReturnValue(true); + // Keep the public spy handle name for existing tests. The deprecated factory + // delegates to createCallerContext, so this spy covers both entry points. const createUserContextSpy = vi - .spyOn(ServiceContext, "createUserContext") + .spyOn(ServiceContext, "createCallerContext") .mockImplementation( fakeUserContext( options.userDatabricksClient || createMockWorkspaceClient(), @@ -545,6 +555,7 @@ export function mockServiceContext(options: TestContextOptions = {}) { initSpy.mockRestore(); isInitializedSpy.mockRestore(); createUserContextSpy.mockRestore(); + resourcesSpy.mockRestore(); }, }; } diff --git a/packages/appkit/src/testing/tests/fixtures.test.ts b/packages/appkit/src/testing/tests/fixtures.test.ts index 9f7a906f8..4382b060e 100644 --- a/packages/appkit/src/testing/tests/fixtures.test.ts +++ b/packages/appkit/src/testing/tests/fixtures.test.ts @@ -4,6 +4,7 @@ import { CacheManager } from "../../cache"; import { InMemoryStorage } from "../../cache/storage"; import { ServiceContext } from "../../context"; import { AuthenticationError } from "../../errors"; +import { getWarehouseId } from "../../resources"; import { ApiError } from "../../workspace-client"; import { createApiError, @@ -154,6 +155,19 @@ describe("mockServiceContext — user context matches production", () => { describe("useServiceContextMock", () => { const ctx = useServiceContextMock({ warehouseId: "wh-1" }); + test("warehouse resources do not depend on the deprecated context field", async () => { + const legacyRead = vi.fn(() => { + throw new Error("Deprecated context warehouse must not be read"); + }); + Object.defineProperty(ctx.current.serviceContext, "warehouseId", { + get: legacyRead, + }); + + expect(await getWarehouseId()).toBe("wh-1"); + expect(legacyRead).not.toHaveBeenCalled(); + expect(ctx.current.getSpy).not.toHaveBeenCalled(); + }); + test(".current exposes the active mock, installed for this test", () => { // The spy is live: the real singleton getter is replaced. expect(vi.isMockFunction(ServiceContext.get)).toBe(true); diff --git a/packages/shared/src/schemas/manifest.test.ts b/packages/shared/src/schemas/manifest.test.ts new file mode 100644 index 000000000..13f2939f8 --- /dev/null +++ b/packages/shared/src/schemas/manifest.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, test } from "vitest"; + +import { + APP_ONLY_RESOURCE_TYPES, + pluginManifestSchema, + SCOPE_BY_TYPE, +} from "./manifest"; + +const manifest = { + name: "testPlugin", + displayName: "Test Plugin", + description: "A test plugin", + resources: { required: [], optional: [] }, +}; + +describe("manifest execution capabilities", () => { + test("only maps resource types with confirmed user scopes", () => { + expect(SCOPE_BY_TYPE).toEqual({ + sql_warehouse: "sql", + serving_endpoint: "model-serving", + genie_space: "genie", + volume: "files", + vector_search_index: "vector-search", + uc_connection: "catalog.connections", + }); + expect([...APP_ONLY_RESOURCE_TYPES]).toEqual([ + "secret", + "database", + "postgres", + ]); + }); + + test("preserves manifests with no authored scopes", () => { + expect(pluginManifestSchema.parse(manifest)).toEqual(manifest); + expect( + pluginManifestSchema.parse({ ...manifest, scopes: [] }).scopes, + ).toEqual([]); + }); + + test("accepts every capability-only scope without resource IDs", () => { + const scopes = [ + "ai-gateway", + "mcp.external", + "mcp.functions", + "workspace.workspace", + "catalog.catalogs:read", + "catalog.schemas:read", + "catalog.tables:read", + ]; + expect(pluginManifestSchema.parse({ ...manifest, scopes }).scopes).toEqual( + scopes, + ); + }); + + test.each([ + "sql", + "sql:restricted-query", + "postgres", + "genie", + "model-serving", + "files", + "vector-search", + "catalog.connections", + "mlflow", + "jobs", + "apps", + "dashboards.genie", + "files.files", + "serving.serving-endpoints", + "unknown", + ])("rejects %s as an authored capability-only scope", (scope) => { + expect( + pluginManifestSchema.safeParse({ ...manifest, scopes: [scope] }).success, + ).toBe(false); + }); + + test("rejects malformed scopes and per-plugin authMode", () => { + expect( + pluginManifestSchema.safeParse({ ...manifest, scopes: "ai-gateway" }) + .success, + ).toBe(false); + expect( + pluginManifestSchema.safeParse({ ...manifest, scopes: [123] }).success, + ).toBe(false); + expect( + pluginManifestSchema.safeParse({ ...manifest, authMode: "obo" }).success, + ).toBe(false); + }); +}); diff --git a/packages/shared/src/schemas/manifest.ts b/packages/shared/src/schemas/manifest.ts index 38bfe51f5..af97f0521 100644 --- a/packages/shared/src/schemas/manifest.ts +++ b/packages/shared/src/schemas/manifest.ts @@ -52,6 +52,39 @@ export const resourceTypeSchema = z ]) .describe("Type of Databricks resource"); +/** Apps user_api_scopes for resource types with confirmed OBO support. */ +export const SCOPE_BY_TYPE = { + sql_warehouse: "sql", // sql:restricted-query is the read-only variant. + serving_endpoint: "model-serving", + genie_space: "genie", + volume: "files", + vector_search_index: "vector-search", + uc_connection: "catalog.connections", + // uc_function uses sql, or mcp.functions through managed MCP. Confirm later. + // experiment and job are SP-only; there is no mlflow or jobs scope. +} as const satisfies Partial>; + +// A postgres user_api_scope exists, so Lakebase is platform-OBO-capable. +// It stays app-only for v1 because the connector connects as the SP today (audit A7). +export const APP_ONLY_RESOURCE_TYPES: ReadonlySet = new Set([ + "secret", + "database", + "postgres", +]); + +/** Capabilities that need a user_api_scope but have no resource ID. */ +export const capabilityScopeSchema = z.enum([ + "ai-gateway", + "mcp.external", + "mcp.functions", + "workspace.workspace", + "catalog.catalogs:read", + "catalog.schemas:read", + "catalog.tables:read", +]); + +export type CapabilityScope = z.infer; + export const secretPermissionSchema = z .enum(["READ", "WRITE", "MANAGE"]) .describe("Permission for secret resources (order: weakest to strongest)"); @@ -668,6 +701,10 @@ export const pluginScaffoldingRulesSchema = z export const pluginManifestSchema = z .object({ + scopes: z + .array(capabilityScopeSchema) + .optional() + .describe("Capability-only user_api_scopes with no resource ID."), $schema: z .string() .optional()