diff --git a/docs/docs/api/appkit/Class.AppKitError.md b/docs/docs/api/appkit/Class.AppKitError.md index 1311a903e..3cb3a0a5b 100644 --- a/docs/docs/api/appkit/Class.AppKitError.md +++ b/docs/docs/api/appkit/Class.AppKitError.md @@ -32,6 +32,7 @@ console.error(error.toJSON()); // Safe for logging, sensitive values redacted - [`ConnectionError`](Class.ConnectionError.md) - [`DatabaseValidationError`](Class.DatabaseValidationError.md) - [`ExecutionError`](Class.ExecutionError.md) +- [`IdentityExpiredError`](Class.IdentityExpiredError.md) - [`InitializationError`](Class.InitializationError.md) - [`ServerError`](Class.ServerError.md) - [`TunnelError`](Class.TunnelError.md) @@ -78,8 +79,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.AuthenticationError.md b/docs/docs/api/appkit/Class.AuthenticationError.md index bc70ceedf..f6842d078 100644 --- a/docs/docs/api/appkit/Class.AuthenticationError.md +++ b/docs/docs/api/appkit/Class.AuthenticationError.md @@ -53,8 +53,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.ConfigurationError.md b/docs/docs/api/appkit/Class.ConfigurationError.md index 5734349d9..2eba5359f 100644 --- a/docs/docs/api/appkit/Class.ConfigurationError.md +++ b/docs/docs/api/appkit/Class.ConfigurationError.md @@ -53,8 +53,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.ConnectionError.md b/docs/docs/api/appkit/Class.ConnectionError.md index afb6e6e3f..c54428a8f 100644 --- a/docs/docs/api/appkit/Class.ConnectionError.md +++ b/docs/docs/api/appkit/Class.ConnectionError.md @@ -53,8 +53,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.DatabaseValidationError.md b/docs/docs/api/appkit/Class.DatabaseValidationError.md index d9315dbf6..d9c49752b 100644 --- a/docs/docs/api/appkit/Class.DatabaseValidationError.md +++ b/docs/docs/api/appkit/Class.DatabaseValidationError.md @@ -40,8 +40,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.ExecutionError.md b/docs/docs/api/appkit/Class.ExecutionError.md index e7c6f4573..359ea3c57 100644 --- a/docs/docs/api/appkit/Class.ExecutionError.md +++ b/docs/docs/api/appkit/Class.ExecutionError.md @@ -55,8 +55,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.IdentityExpiredError.md b/docs/docs/api/appkit/Class.IdentityExpiredError.md new file mode 100644 index 000000000..4e2f3e3bb --- /dev/null +++ b/docs/docs/api/appkit/Class.IdentityExpiredError.md @@ -0,0 +1,187 @@ +# Class: IdentityExpiredError + +The downstream service rejected the active caller's credentials. + +## Extends + +- [`AppKitError`](Class.AppKitError.md) + +## Constructors + +### Constructor + +```ts +new IdentityExpiredError(tokenFingerprint?: string): IdentityExpiredError; +``` + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `tokenFingerprint?` | `string` | + +#### Returns + +`IdentityExpiredError` + +#### Overrides + +[`AppKitError`](Class.AppKitError.md).[`constructor`](Class.AppKitError.md#constructor) + +## Properties + +### \_clientMessage? + +```ts +protected readonly optional _clientMessage: string; +``` + +Client-safe error message. When set, callers serializing the error to +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement +fragments, internal object names, and correlation IDs. + +Subclasses can set this in their constructor for a fixed sanitized +string. When unset, `clientMessage` defaults to a generic per-code +string (see the getter), and the raw `message` is kept server-side +only. + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`_clientMessage`](Class.AppKitError.md#_clientmessage) + +*** + +### cause? + +```ts +readonly optional cause: Error; +``` + +Optional cause of the error + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`cause`](Class.AppKitError.md#cause) + +*** + +### code + +```ts +readonly code: "IDENTITY_EXPIRED" = "IDENTITY_EXPIRED"; +``` + +Error code for programmatic error handling + +#### Overrides + +[`AppKitError`](Class.AppKitError.md).[`code`](Class.AppKitError.md#code) + +*** + +### context? + +```ts +readonly optional context: Record; +``` + +Additional context for the error + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`context`](Class.AppKitError.md#context) + +*** + +### isRetryable + +```ts +readonly isRetryable: false = false; +``` + +Whether this error type is generally safe to retry + +#### Overrides + +[`AppKitError`](Class.AppKitError.md).[`isRetryable`](Class.AppKitError.md#isretryable) + +*** + +### statusCode + +```ts +readonly statusCode: 401 = 401; +``` + +HTTP status code suggestion (can be overridden) + +#### Overrides + +[`AppKitError`](Class.AppKitError.md).[`statusCode`](Class.AppKitError.md#statuscode) + +*** + +### tokenFingerprint? + +```ts +readonly optional tokenFingerprint: string; +``` + +## Accessors + +### clientMessage + +#### Get Signature + +```ts +get clientMessage(): string; +``` + +Sanitized message safe to forward to clients. Override in subclasses +if a more specific default is appropriate. + +##### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`clientMessage`](Class.AppKitError.md#clientmessage) + +## Methods + +### toJSON() + +```ts +toJSON(): Record; +``` + +Convert error to JSON for logging/serialization. +Sensitive values in context are automatically redacted. + +#### Returns + +`Record`\<`string`, `unknown`\> + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`toJSON`](Class.AppKitError.md#tojson) + +*** + +### toString() + +```ts +toString(): string; +``` + +Create a human-readable string representation + +#### Returns + +`string` + +#### Inherited from + +[`AppKitError`](Class.AppKitError.md).[`toString`](Class.AppKitError.md#tostring) diff --git a/docs/docs/api/appkit/Class.InitializationError.md b/docs/docs/api/appkit/Class.InitializationError.md index 0bc119431..aa3327be6 100644 --- a/docs/docs/api/appkit/Class.InitializationError.md +++ b/docs/docs/api/appkit/Class.InitializationError.md @@ -53,8 +53,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.Plugin.md b/docs/docs/api/appkit/Class.Plugin.md index a783c2c92..1d6b31fb6 100644 --- a/docs/docs/api/appkit/Class.Plugin.md +++ b/docs/docs/api/appkit/Class.Plugin.md @@ -364,7 +364,8 @@ Returns an [ExecutionResult](TypeAlias.ExecutionResult.md) discriminated union: - `{ ok: true, data: T }` on success - `{ ok: false, status: number, message: string }` on failure -Errors are never thrown — the method is production-safe. +Caller credential expiration retains the failure result and additionally +exposes a typed error, preserving existing result-based callers. #### Type Parameters diff --git a/docs/docs/api/appkit/Class.ServerError.md b/docs/docs/api/appkit/Class.ServerError.md index d2b61182c..1c832dc36 100644 --- a/docs/docs/api/appkit/Class.ServerError.md +++ b/docs/docs/api/appkit/Class.ServerError.md @@ -52,8 +52,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.TunnelError.md b/docs/docs/api/appkit/Class.TunnelError.md index bfea22bd2..ab5bb3ad2 100644 --- a/docs/docs/api/appkit/Class.TunnelError.md +++ b/docs/docs/api/appkit/Class.TunnelError.md @@ -53,8 +53,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/Class.ValidationError.md b/docs/docs/api/appkit/Class.ValidationError.md index f1c760cb1..56d2e7a5e 100644 --- a/docs/docs/api/appkit/Class.ValidationError.md +++ b/docs/docs/api/appkit/Class.ValidationError.md @@ -53,8 +53,8 @@ protected readonly optional _clientMessage: string; ``` Client-safe error message. When set, callers serializing the error to -a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` -— `message` may contain raw upstream / SDK text including statement +a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. +`message` may contain raw upstream / SDK text including statement fragments, internal object names, and correlation IDs. Subclasses can set this in their constructor for a fixed sanitized diff --git a/docs/docs/api/appkit/TypeAlias.ExecutionResult.md b/docs/docs/api/appkit/TypeAlias.ExecutionResult.md index 1af9686e8..bcd581ce8 100644 --- a/docs/docs/api/appkit/TypeAlias.ExecutionResult.md +++ b/docs/docs/api/appkit/TypeAlias.ExecutionResult.md @@ -7,6 +7,7 @@ type ExecutionResult = ok: true; } | { + error?: IdentityExpiredError; message: string; ok: false; status: number; @@ -31,3 +32,59 @@ In production, error messages from non-AppKitError sources are handled as: | Type Parameter | | ------ | | `T` | + +## Type Declaration + +```ts +{ + data: T; + ok: true; +} +``` + +### data + +```ts +data: T; +``` + +### ok + +```ts +ok: true; +``` + +```ts +{ + error?: IdentityExpiredError; + message: string; + ok: false; + status: number; +} +``` + +### error? + +```ts +optional error: IdentityExpiredError; +``` + +Typed credential expiry without changing the existing failure envelope. + +### message + +```ts +message: string; +``` + +### ok + +```ts +ok: false; +``` + +### status + +```ts +status: number; +``` diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index dd6af49b8..f32f79917 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -22,6 +22,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [DatabaseValidationError](Class.DatabaseValidationError.md) | Deliberate validation failure raised by a database mutation hook. Generated routes answer `422` and echo only the issues naming a public column; every other failure raised inside a hook stays an opaque server error. | | [DatabricksAdapter](Class.DatabricksAdapter.md) | Adapter that talks directly to Databricks Model Serving `/invocations` endpoint. | | [ExecutionError](Class.ExecutionError.md) | Error thrown when an operation execution fails. Use for statement failures, canceled operations, or unexpected states. | +| [IdentityExpiredError](Class.IdentityExpiredError.md) | The downstream service rejected the active caller's credentials. | | [InitializationError](Class.InitializationError.md) | Error thrown when a service or component is not properly initialized. Use when accessing services before they are ready. | | [MlflowClient](Class.MlflowClient.md) | A thin client over the Databricks workspace REST API, owning the host + bearer token so callers (eval-run creation, assessment writes, the judge's serving endpoint) don't each re-derive URLs or re-attach auth. The host is normalized once at construction. | | [Plugin](Class.Plugin.md) | Base abstract class for creating AppKit plugins. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 0378d3e60..176a322cb 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -61,6 +61,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Class.ExecutionError", label: "ExecutionError" }, + { + type: "doc", + id: "api/appkit/Class.IdentityExpiredError", + label: "IdentityExpiredError" + }, { type: "doc", id: "api/appkit/Class.InitializationError", diff --git a/docs/docs/plugins/execution-context.md b/docs/docs/plugins/execution-context.md index 94ab16b18..1764bf7c5 100644 --- a/docs/docs/plugins/execution-context.md +++ b/docs/docs/plugins/execution-context.md @@ -113,3 +113,64 @@ runs with default app credentials, marked `DEV_OBO_FALLBACK`. If a caller scope is already open, fallback retains it instead of widening to SP. The marker does not leak outside the scope. Production never falls back when credentials are missing. + +## Credential expiration and telemetry + +A structured downstream HTTP 401 inside a caller scope throws +`IdentityExpiredError` with code `IDENTITY_EXPIRED`. It includes the existing +token fingerprint, not the token or upstream credential-bearing error. Obtain +fresh user credentials before retrying. Non-401 failures and SP execution keep +their existing behavior. Plugin `execute()` preserves its failed-result envelope +and adds the typed error in the optional `error` field; SSE streams expose `IDENTITY_EXPIRED` +in the error payload's `errorCode` field. + +AppKit-managed spans include `appkit.execution.principal` (`app` or `user`) and +`appkit.execution.principal_id` (user or SP ID, or `app` before initialization). +`appkit.execution.actor_id` is present when an initiating user exists. Tokens +are never attached to these attributes. + +## Real user execution locally + +Set `DATABRICKS_TOKEN` and `DATABRICKS_HOST` in the app's `.env` to use a user +token directly: + +```dotenv +DATABRICKS_HOST=https://your-workspace.cloud.databricks.com +DATABRICKS_TOKEN=your-user-token +``` + +When `DATABRICKS_TOKEN` is present, it takes precedence. AppKit uses it directly +and resolves the user ID from the configured host. Otherwise, set +`DATABRICKS_CONFIG_PROFILE` to an authenticated user profile for the same +workspace as the app. The generated template already sets the profile when you +choose one during scaffolding: + +```dotenv +DATABRICKS_CONFIG_PROFILE=your-user-profile +``` + +Then run your usual command: + +```sh +npm run dev +``` + +Open the app's normal localhost URL. In development, the server automatically +adds `x-forwarded-access-token`, `x-forwarded-user`, and optional email headers +before plugin routes and custom routes run. No separate proxy, target, or port +is needed. `asUser(req)` uses that user identity. Unscoped operations still use +the app's configured credentials; injecting headers does not open a caller scope. +For a genuine SP-versus-user comparison, the app credentials must belong to an +SP, not the same user profile. + +Credentials stay in memory and refresh after 30 seconds of use. Initial auth and +refresh failures return 401, never a silent SP fallback. Existing forwarded user +tokens are preserved. Automatic injection runs only in `NODE_ENV=development` +and only for same-origin loopback requests, including `localhost`. Use it only +with a trusted local app. It emulates user credentials, not platform consent, +scope enforcement, or resource provisioning. + +Set `APPKIT_DEV_OBO=false` in `.env` to disable automatic injection. Without a +configured token or profile, injection is also disabled. Tokenless +`asUser(req)` then keeps the existing `DEV_OBO_FALLBACK` behavior in +development. diff --git a/packages/appkit/src/context/execution-context.ts b/packages/appkit/src/context/execution-context.ts index dc5ea6bd4..9f87ddbf4 100644 --- a/packages/appkit/src/context/execution-context.ts +++ b/packages/appkit/src/context/execution-context.ts @@ -1,5 +1,9 @@ import { AsyncLocalStorage } from "node:async_hooks"; +import { + IdentityExpiredError, + isUnauthorized, +} from "../errors/identity-expired"; import { captureWarehouseId, getWarehouseId as getResourceWarehouseId, @@ -24,10 +28,34 @@ function runInCallerScope( ): T { const caller = snapshotCallerContext(callerContext); return runWithResourceBindings(legacyResources, () => - executionContextStorage.run(caller, fn), + executionContextStorage.run(caller, () => { + try { + const result = fn(); + if (result instanceof Promise) { + return result.catch((error) => { + throw normalizeIdentityError(error); + }) as T; + } + return result; + } catch (error) { + throw normalizeIdentityError(error); + } + }), ); } +/** Convert downstream 401s before recording or surfacing caller failures. */ +export function normalizeIdentityError(error: unknown): unknown { + const caller = getCallerContext(); + if ( + !caller || + error instanceof IdentityExpiredError || + !isUnauthorized(error) + ) + return error; + return new IdentityExpiredError(caller.tokenFingerprint); +} + /** * Run a function with an immutable snapshot of the caller context. * Nested and concurrent scopes keep their own identities. diff --git a/packages/appkit/src/context/service-context.ts b/packages/appkit/src/context/service-context.ts index 0475e8a95..1d463a06b 100644 --- a/packages/appkit/src/context/service-context.ts +++ b/packages/appkit/src/context/service-context.ts @@ -111,7 +111,12 @@ export class ServiceContext { throw AuthenticationError.missingToken("user token"); } - const host = process.env.DATABRICKS_HOST; + // Local templates can configure only a profile, whose host the SDK resolved. + const host = + process.env.DATABRICKS_HOST || + (process.env.NODE_ENV === "development" && ServiceContext.isInitialized() + ? ServiceContext.get().client.config?.host + : undefined); if (!host) { throw ConfigurationError.missingEnvVar("DATABRICKS_HOST"); } diff --git a/packages/appkit/src/context/tests/identity-lifecycle.test.ts b/packages/appkit/src/context/tests/identity-lifecycle.test.ts new file mode 100644 index 000000000..b1fb6fc08 --- /dev/null +++ b/packages/appkit/src/context/tests/identity-lifecycle.test.ts @@ -0,0 +1,205 @@ +import type { AgentToolDefinition, ToolProvider } from "shared"; +import { describe, expect, test } from "vitest"; + +import { IdentityExpiredError } from "../../errors"; +import { AuthenticationError } from "../../errors/authentication"; +import { Plugin, toPlugin } from "../../plugin"; +import { agents } from "../../plugins/agents"; +import { + createMockRequest, + createMockWorkspaceClient, + createTestApp, + createTestPluginContext, +} from "../../testing"; +import { + getCallerContext, + normalizeIdentityError, + runInCallerContext, +} from "../execution-context"; + +const caller = { + principal: { type: "user" as const, userId: "alice" }, + client: createMockWorkspaceClient(), + workspaceId: Promise.resolve("workspace"), + tokenFingerprint: "1234567890abcdef", +}; +const upstream = Object.assign(new Error("secret-bearer-token"), { + statusCode: 401, +}); + +class ExpiringToolsPlugin extends Plugin implements ToolProvider { + static manifest = { + name: "expiringTools" as const, + displayName: "Expiring tools", + description: "probe", + resources: { required: [], optional: [] }, + }; + getAgentTools(): AgentToolDefinition[] { + return [ + { + name: "fail", + description: "Fails with an expired user token", + parameters: { type: "object", properties: {} }, + }, + ]; + } + async executeAgentTool(): Promise { + throw upstream; + } +} +const expiringTools = toPlugin(ExpiringToolsPlugin); +class ExpiringPlugin extends Plugin { + static manifest = { + name: "expiring" as const, + displayName: "Expiring", + description: "probe", + resources: { required: [], optional: [] }, + }; + exports() { + return { + query: () => + this.execute( + async () => { + throw upstream; + }, + { default: {} }, + ), + stream: async function* () { + yield "start"; + throw upstream; + }, + }; + } +} + +describe("identity expiration", () => { + test("agent HTTP responses expose a safe identity-expiration code", async () => { + await using app = await createTestApp({ + plugins: [ + expiringTools(), + agents({ + agents: { + probe: { + instructions: "probe", + // The user token expires during a plugin tool call, which runs + // in user scope; the model call itself runs as the app. + model: { + async *run(_input, ctx) { + yield { type: "message_delta", content: "" }; + await ctx.executeTool("expiringTools.fail", {}); + }, + }, + tools: (plugins) => plugins.expiringTools.toolkit(), + }, + }, + }), + ], + }); + const response = await app.post("/invocations", { + body: { input: "hi", stream: false }, + obo: { userId: "alice" }, + }); + expect(response.status).toBe(401); + expect(await response.json()).toMatchObject({ code: "IDENTITY_EXPIRED" }); + const streamed = await app.post("/api/agents/chat", { + body: { message: "hi" }, + obo: { userId: "alice" }, + }); + const body = await streamed.text(); + expect(body).toContain("IDENTITY_EXPIRED"); + expect(body).not.toContain("secret-bearer-token"); + }); + test.each([ + { status: 401 }, + { response: { status: 401 } }, + new Error("wrapper", { cause: upstream }), + ])( + "converts structured downstream failures without retaining credentials", + async (error) => { + const result = runInCallerContext(caller, async () => { + throw error; + }); + const expired = await result.catch((err) => err); + expect(expired).toBeInstanceOf(IdentityExpiredError); + expect(expired.tokenFingerprint).toBe(caller.tokenFingerprint); + expect(expired.isRetryable).toBe(false); + expect(expired.statusCode).toBe(401); + expect(expired.cause).toBeUndefined(); + expect(JSON.stringify(expired)).not.toContain("secret-bearer-token"); + expect(String(expired)).not.toContain("secret-bearer-token"); + expect(getCallerContext()).toBeUndefined(); + }, + ); + + test("preserves non-401, local missing-token errors, and already typed failures", () => { + for (const error of [ + new Error("401 in text"), + { statusCode: 403 }, + AuthenticationError.missingToken(), + new IdentityExpiredError("fingerprint"), + ]) { + try { + runInCallerContext(caller, () => { + throw error; + }); + } catch (result) { + expect(result).toBe(error); + } + } + }); + + test("normalizes tool failures under the caller before recording telemetry", async () => { + const mock = createTestPluginContext({ + provider: { + tool: () => { + throw upstream; + }, + }, + }); + await expect( + runInCallerContext(caller, () => + mock.ctx.executeTool( + createMockRequest() as never, + "provider", + "tool", + {}, + ), + ), + ).rejects.toBeInstanceOf(IdentityExpiredError); + await expect( + mock.ctx.executeTool( + createMockRequest({ obo: { userId: "alice" } }) as never, + "provider", + "tool", + {}, + ), + ).rejects.toBeInstanceOf(IdentityExpiredError); + expect(normalizeIdentityError(upstream)).toBe(upstream); + }); + + test("plugin execution and delayed stream iteration surface the typed error", async () => { + await using app = await createTestApp({ + plugins: [toPlugin(ExpiringPlugin)()], + server: false, + }); + const scoped = app.plugins.asUser( + createMockRequest({ obo: { userId: "alice" } }), + ); + await expect(scoped.expiring.query()).resolves.toMatchObject({ + ok: false, + status: 401, + error: expect.any(IdentityExpiredError), + }); + await expect( + app.plugins.expiring.asUser(createMockRequest({ obo: true })).query(), + ).resolves.toMatchObject({ ok: false, status: 401 }); + const stream = scoped.expiring.stream(); + expect(await stream.next()).toEqual({ value: "start", done: false }); + await expect(stream.next()).rejects.toBeInstanceOf(IdentityExpiredError); + expect(getCallerContext()).toBeUndefined(); + expect(await app.plugins.expiring.query()).toMatchObject({ + ok: false, + status: 401, + }); + }); +}); diff --git a/packages/appkit/src/context/tests/service-context.test.ts b/packages/appkit/src/context/tests/service-context.test.ts index 470132142..4bb5f27f6 100644 --- a/packages/appkit/src/context/tests/service-context.test.ts +++ b/packages/appkit/src/context/tests/service-context.test.ts @@ -317,6 +317,26 @@ describe("ServiceContext", () => { ); }); + test("uses the initialized profile host for local callers without DATABRICKS_HOST", () => { + delete process.env.DATABRICKS_HOST; + process.env.NODE_ENV = "development"; + process.env.DATABRICKS_CONFIG_PROFILE = "selected-user"; + Object.defineProperty(ServiceContext.get().client, "config", { + value: { host: "https://profile-workspace.databricks.com" }, + }); + + const caller = ServiceContext.createCallerContext("user-token", "alice"); + + expect(caller.principal).toMatchObject({ type: "user", userId: "alice" }); + expect(MockWorkspaceClient).toHaveBeenLastCalledWith( + expect.objectContaining({ + host: "https://profile-workspace.databricks.com", + token: "user-token", + authType: "pat", + }), + ); + }); + test("should throw InitializationError when service context is not initialized", () => { ServiceContext.reset(); diff --git a/packages/appkit/src/core/plugin-context.ts b/packages/appkit/src/core/plugin-context.ts index 173fb965f..6605d8161 100644 --- a/packages/appkit/src/core/plugin-context.ts +++ b/packages/appkit/src/core/plugin-context.ts @@ -1,7 +1,10 @@ import type express from "express"; import type { BasePlugin, ToolProvider } from "shared"; -import { getCallerContext } from "../context/execution-context"; +import { + getCallerContext, + normalizeIdentityError, +} from "../context/execution-context"; import { createRequestScope } from "../context/request-scope"; import { ServiceContext } from "../context/service-context"; import { createLogger } from "../logging/logger"; @@ -246,7 +249,8 @@ export class PluginContext { ); span.setStatus({ code: SpanStatusCode.OK }); return result; - } catch (error) { + } catch (caught) { + const error = normalizeIdentityError(caught); span.setStatus({ code: SpanStatusCode.ERROR, message: diff --git a/packages/appkit/src/errors/base.ts b/packages/appkit/src/errors/base.ts index 7ed56d213..1732cd9e5 100644 --- a/packages/appkit/src/errors/base.ts +++ b/packages/appkit/src/errors/base.ts @@ -48,8 +48,8 @@ export abstract class AppKitError extends Error { /** * Client-safe error message. When set, callers serializing the error to - * a client (SSE, HTTP body) MUST prefer `clientMessage` over `message` - * — `message` may contain raw upstream / SDK text including statement + * a client (SSE, HTTP body) MUST prefer `clientMessage` over `message`. + * `message` may contain raw upstream / SDK text including statement * fragments, internal object names, and correlation IDs. * * Subclasses can set this in their constructor for a fixed sanitized diff --git a/packages/appkit/src/errors/identity-expired.ts b/packages/appkit/src/errors/identity-expired.ts new file mode 100644 index 000000000..9cfa7f1ff --- /dev/null +++ b/packages/appkit/src/errors/identity-expired.ts @@ -0,0 +1,42 @@ +import { AppKitError } from "./base"; + +/** The downstream service rejected the active caller's credentials. */ +export class IdentityExpiredError extends AppKitError { + readonly code = "IDENTITY_EXPIRED"; + readonly statusCode = 401; + readonly isRetryable = false; + + constructor(readonly tokenFingerprint?: string) { + const message = + "Caller credentials were rejected or expired. Reauthenticate and retry with a fresh user token."; + // Do not retain upstream errors: SDK errors may contain authorization headers. + super(message, { + clientMessage: message, + context: { fingerprint: tokenFingerprint }, + }); + } +} + +/** Recognize structured HTTP failures, including connector cause wrappers. */ +export function isUnauthorized(error: unknown): boolean { + const seen = new Set(); + let current = error; + while (current && typeof current === "object" && !seen.has(current)) { + seen.add(current); + const value = current as { + status?: unknown; + statusCode?: unknown; + response?: { status?: unknown }; + cause?: unknown; + }; + if ( + !(current instanceof AppKitError) && + (value.status === 401 || + value.statusCode === 401 || + value.response?.status === 401) + ) + return true; + current = value.cause; + } + return false; +} diff --git a/packages/appkit/src/errors/index.ts b/packages/appkit/src/errors/index.ts index 22ab10f0c..24dbb7f8f 100644 --- a/packages/appkit/src/errors/index.ts +++ b/packages/appkit/src/errors/index.ts @@ -20,6 +20,7 @@ */ export { AuthenticationError } from "./authentication"; +export { IdentityExpiredError } from "./identity-expired"; export { AppKitError } from "./base"; export { ConfigurationError } from "./configuration"; export { ConnectionError } from "./connection"; diff --git a/packages/appkit/src/index.ts b/packages/appkit/src/index.ts index 650e875a9..6bc5f1b28 100644 --- a/packages/appkit/src/index.ts +++ b/packages/appkit/src/index.ts @@ -68,6 +68,7 @@ export { DatabaseValidationError, type DatabaseValidationIssue, ExecutionError, + IdentityExpiredError, InitializationError, ServerError, TunnelError, diff --git a/packages/appkit/src/plugin/execution-result.ts b/packages/appkit/src/plugin/execution-result.ts index 964e89839..25bb32b69 100644 --- a/packages/appkit/src/plugin/execution-result.ts +++ b/packages/appkit/src/plugin/execution-result.ts @@ -1,3 +1,5 @@ +import type { IdentityExpiredError } from "../errors/identity-expired"; + /** * Discriminated union for plugin execution results. * @@ -14,4 +16,10 @@ */ export type ExecutionResult = | { ok: true; data: T } - | { ok: false; status: number; message: string }; + | { + ok: false; + status: number; + message: string; + /** Typed credential expiry without changing the existing failure envelope. */ + error?: IdentityExpiredError; + }; diff --git a/packages/appkit/src/plugin/interceptors/telemetry.ts b/packages/appkit/src/plugin/interceptors/telemetry.ts index f21bb5517..4759897ff 100644 --- a/packages/appkit/src/plugin/interceptors/telemetry.ts +++ b/packages/appkit/src/plugin/interceptors/telemetry.ts @@ -3,6 +3,7 @@ import type { TelemetryConfig } from "shared"; import { getCurrentUserId, isInUserContext, + normalizeIdentityError, } from "../../context/execution-context"; import type { ITelemetry, Span } from "../../telemetry"; import { SpanStatusCode } from "../../telemetry"; @@ -65,7 +66,8 @@ export class TelemetryInterceptor implements ExecutionInterceptor { span.setStatus({ code: SpanStatusCode.OK }); } return result; - } catch (error) { + } catch (caught) { + const error = normalizeIdentityError(caught); if (!isAborted) { span.recordException(error as Error); span.setStatus({ code: SpanStatusCode.ERROR }); diff --git a/packages/appkit/src/plugin/plugin.ts b/packages/appkit/src/plugin/plugin.ts index 474284c1b..3b0d7418b 100644 --- a/packages/appkit/src/plugin/plugin.ts +++ b/packages/appkit/src/plugin/plugin.ts @@ -18,10 +18,12 @@ import { AppManager } from "../app"; import { CacheManager } from "../cache"; import { getCurrentUserId } from "../context"; import { warnContextDeprecation } from "../context/deprecation"; +import { normalizeIdentityError } from "../context/execution-context"; import { createRequestScope } from "../context/request-scope"; import { scopePlugin } from "../context/scoped-api"; import type { PluginContext } from "../core/plugin-context"; import { AppKitError, AuthenticationError } from "../errors"; +import { IdentityExpiredError } from "../errors/identity-expired"; import { createLogger } from "../logging/logger"; import { StreamManager } from "../stream"; import { @@ -409,10 +411,14 @@ export abstract class Plugin< ); // check if result is a generator - if (self._checkIfGenerator(result)) { - yield* result; - } else { - yield result; + try { + if (self._checkIfGenerator(result)) { + yield* result; + } else { + yield result; + } + } catch (error) { + throw normalizeIdentityError(error); } }; @@ -435,7 +441,8 @@ export abstract class Plugin< * - `{ ok: true, data: T }` on success * - `{ ok: false, status: number, message: string }` on failure * - * Errors are never thrown — the method is production-safe. + * Caller credential expiration retains the failure result and additionally + * exposes a typed error, preserving existing result-based callers. */ protected async execute( fn: (signal?: AbortSignal) => Promise, @@ -461,7 +468,8 @@ export abstract class Plugin< context, ); return { ok: true, data }; - } catch (error) { + } catch (caught) { + const error = normalizeIdentityError(caught); logger.error("Plugin execution failed", { error, plugin: this.name }); if (error instanceof AppKitError) { @@ -469,6 +477,7 @@ export abstract class Plugin< ok: false, status: error.statusCode, message: error.message, + ...(error instanceof IdentityExpiredError ? { error } : {}), }; } diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index 059136cb9..6856b2907 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -19,6 +19,7 @@ import type { import { isSupervisorTool } from "../../agents/supervisor-api"; import { AppKitMcpClient, buildMcpHostPolicy } from "../../connectors/mcp"; import { getWorkspaceClient } from "../../context"; +import { normalizeIdentityError } from "../../context/execution-context"; import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { loadAgentsFromDir } from "../../core/agent/load-agents"; import { CODE_AGENTS_SOURCE_DIR } from "../../core/agent/load-code-agents"; @@ -44,6 +45,7 @@ import type { ResolvedToolEntry, } from "../../core/agent/types"; import { isToolkitEntry } from "../../core/agent/types"; +import { IdentityExpiredError } from "../../errors/identity-expired"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; @@ -1349,8 +1351,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { outboundEvents.close(); return; } - logger.error("Agent chat error: %O", error); - outboundEvents.close(error); + const failure = normalizeIdentityError(error); + logger.error("Agent chat error: %O", failure); + outboundEvents.close(failure); return; } finally { // Any pending approval gates for this stream are auto-denied so the @@ -1531,7 +1534,14 @@ export class AgentsPlugin extends Plugin implements ToolProvider { res.status(499).json({ error: "Request aborted" }); return; } - logger.error("Agent invoke error: %O", error); + const failure = normalizeIdentityError(error); + if (failure instanceof IdentityExpiredError) { + res + .status(401) + .json({ error: failure.clientMessage, code: failure.code }); + return; + } + logger.error("Agent invoke error: %O", failure); const message = process.env.NODE_ENV === "production" ? "Internal server error" diff --git a/packages/appkit/src/plugins/agents/mlflow.ts b/packages/appkit/src/plugins/agents/mlflow.ts index f73b5a7ea..ac985241c 100644 --- a/packages/appkit/src/plugins/agents/mlflow.ts +++ b/packages/appkit/src/plugins/agents/mlflow.ts @@ -2,6 +2,7 @@ import type { UnityCatalogLocation } from "@mlflow/core"; import { SpanKind } from "@opentelemetry/api"; import type { SpanProcessor } from "@opentelemetry/sdk-trace-base"; +import { normalizeIdentityError } from "../../context/execution-context"; import { createLogger } from "../../logging/logger"; import { TelemetryManager } from "../../telemetry"; @@ -480,15 +481,22 @@ async function trace( inputs: unknown, fn: (span: SpanRecorder) => Promise, ): Promise { - if (!enabled || !mlflow) return fn(noopRecorder); + const execute = async (span: SpanRecorder): Promise => { + try { + return await fn(span); + } catch (error) { + throw normalizeIdentityError(error); + } + }; + if (!enabled || !mlflow) return execute(noopRecorder); const m = mlflow; - if (!ensureConfigured()) return fn(noopRecorder); + if (!ensureConfigured()) return execute(noopRecorder); const type = spanType === "AGENT" ? m.SpanType.AGENT : m.SpanType.TOOL; return await m.withSpan( async (span) => { if (inputs !== undefined) span.setInputs(inputs); let outputsSet = false; - const result = await fn({ + const result = await execute({ setOutputs(outputs) { outputsSet = true; span.setOutputs(outputs); diff --git a/packages/appkit/src/plugins/agents/tool-dispatch.ts b/packages/appkit/src/plugins/agents/tool-dispatch.ts index 9b73772f5..492209341 100644 --- a/packages/appkit/src/plugins/agents/tool-dispatch.ts +++ b/packages/appkit/src/plugins/agents/tool-dispatch.ts @@ -4,6 +4,7 @@ import type express from "express"; import type { AgentRunContext, Message, ResponseStreamEvent } from "shared"; import type { AppKitMcpClient } from "../../connectors/mcp"; +import { normalizeIdentityError } from "../../context/execution-context"; import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream"; import { normalizeToolResult } from "../../core/agent/normalize-result"; import type { @@ -142,7 +143,8 @@ export async function dispatchToolCall( toolResult = await traceTool(name, args, () => runToolEntry(deps, runState, entry, name, args, depth), ); - } catch (err) { + } catch (caught) { + const err = normalizeIdentityError(caught); const error = err instanceof Error ? err.message : String(err); logger.error( "Tool '%s' failed (request %s): %O", diff --git a/packages/appkit/src/plugins/server/dev-obo-middleware.ts b/packages/appkit/src/plugins/server/dev-obo-middleware.ts new file mode 100644 index 000000000..4647ba0b2 --- /dev/null +++ b/packages/appkit/src/plugins/server/dev-obo-middleware.ts @@ -0,0 +1,80 @@ +import type { Request, RequestHandler } from "express"; +import { + createDevOboIdentityProvider, + loadDevOboIdentityFromEnvironment, +} from "shared"; + +import { createLogger } from "../../logging/logger"; + +const logger = createLogger("server:dev-obo"); +const loopbackHosts = new Set(["localhost", "127.0.0.1", "[::1]"]); +const loopbackPeers = new Set(["127.0.0.1", "::1", "::ffff:127.0.0.1"]); + +function isLocalRequest(req: Request): boolean { + // Check the socket, not Express's proxy-aware IP or forwarded headers. + if (!loopbackPeers.has(req.socket.remoteAddress ?? "")) return false; + try { + const origin = new URL(`http://${req.headers.host}`); + if ( + !loopbackHosts.has(origin.hostname) || + origin.host !== req.headers.host || + Number(origin.port || 80) !== req.socket.localPort + ) + return false; + if (req.headers.origin && req.headers.origin !== origin.origin) + return false; + const site = req.headers["sec-fetch-site"]; + return !site || site === "same-origin" || site === "none"; + } catch { + return false; + } +} + +/** Inject local user headers without changing the default execution principal. */ +export function createDevOboMiddleware(): RequestHandler | undefined { + const token = process.env.DATABRICKS_TOKEN?.trim(); + const profile = process.env.DATABRICKS_CONFIG_PROFILE?.trim(); + if ( + process.env.NODE_ENV !== "development" || + process.env.APPKIT_DEV_OBO === "false" || + (!token && !profile) + ) + return undefined; + + const getIdentity = createDevOboIdentityProvider(() => + loadDevOboIdentityFromEnvironment(), + ); + logger.info( + `Local OBO headers enabled using ${token ? "DATABRICKS_TOKEN" : "DATABRICKS_CONFIG_PROFILE"}. Open the app on localhost. Set APPKIT_DEV_OBO=false to disable.`, + ); + + return async (req, res, next) => { + // Preserve explicit forwarded identities. + if (req.header("x-forwarded-access-token")?.trim()) { + next(); + return; + } + if (!isLocalRequest(req)) { + res.status(403).json({ + error: + "Local OBO emulation only accepts same-origin loopback requests.", + }); + return; + } + try { + const identity = await getIdentity(); + req.headers["x-forwarded-access-token"] = identity.token; + req.headers["x-forwarded-user"] = identity.userId; + delete req.headers["x-forwarded-email"]; + if (identity.email) req.headers["x-forwarded-email"] = identity.email; + } catch { + // Never fall back to SP or expose credential-bearing CLI errors. + res.status(401).json({ + error: + "Local OBO credentials unavailable. Verify DATABRICKS_TOKEN and DATABRICKS_HOST, or authenticate the user profile in DATABRICKS_CONFIG_PROFILE, then retry.", + }); + return; + } + next(); + }; +} diff --git a/packages/appkit/src/plugins/server/index.ts b/packages/appkit/src/plugins/server/index.ts index bcdea787a..19951a040 100644 --- a/packages/appkit/src/plugins/server/index.ts +++ b/packages/appkit/src/plugins/server/index.ts @@ -15,6 +15,7 @@ import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; import { instrumentations } from "../../telemetry"; import { sanitizeClientConfig } from "./client-config-sanitizer"; +import { createDevOboMiddleware } from "./dev-obo-middleware"; import manifest from "./manifest.json"; import { RemoteTunnelController } from "./remote-tunnel/remote-tunnel-controller"; import { StaticServer } from "./static-server"; @@ -128,6 +129,8 @@ export class ServerPlugin extends Plugin { */ async start(): Promise { this.serverApplication.use(requestMetricsMiddleware); + const devObo = createDevOboMiddleware(); + if (devObo) this.serverApplication.use(devObo); this.serverApplication.use( express.json({ // Express's stock 100kb default is too tight for modern apps — diff --git a/packages/appkit/src/plugins/server/tests/dev-obo-middleware.test.ts b/packages/appkit/src/plugins/server/tests/dev-obo-middleware.test.ts new file mode 100644 index 000000000..5f71f7bfc --- /dev/null +++ b/packages/appkit/src/plugins/server/tests/dev-obo-middleware.test.ts @@ -0,0 +1,182 @@ +import { Socket } from "node:net"; + +import { request, response, type Response } from "express"; +import { loadDevOboIdentityFromEnvironment } from "shared"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { createDevOboMiddleware } from "../dev-obo-middleware"; + +vi.mock("shared", async (importOriginal) => ({ + ...(await importOriginal()), + loadDevOboIdentityFromEnvironment: vi.fn(), +})); + +function makeRequest() { + const req = Object.create(request) as typeof request; + req.headers = { host: "localhost:8000" }; + req.socket = new Socket(); + Object.defineProperties(req.socket, { + remoteAddress: { value: "127.0.0.1", configurable: true }, + localPort: { value: 8000 }, + }); + return req; +} + +async function invoke(req = makeRequest()) { + const res = Object.create(response) as Response; + const json = vi.fn(() => res); + const status = vi.fn(() => res); + res.json = json; + res.status = status; + const next = vi.fn(); + await createDevOboMiddleware()?.(req, res, next); + return { req, status, json, next }; +} + +beforeEach(() => { + vi.stubEnv("NODE_ENV", "development"); + vi.stubEnv("DATABRICKS_CONFIG_PROFILE", "selected-user"); + vi.stubEnv("DATABRICKS_HOST", undefined); + vi.stubEnv("DATABRICKS_TOKEN", undefined); + vi.stubEnv("APPKIT_DEV_OBO", undefined); + vi.mocked(loadDevOboIdentityFromEnvironment).mockResolvedValue({ + token: "fake-token", + userId: "alice", + email: "alice@example.com", + }); +}); + +afterEach(() => { + vi.unstubAllEnvs(); + vi.clearAllMocks(); +}); + +describe("automatic dev OBO middleware", () => { + test("uses the configured profile and injects the complete caller identity", async () => { + // Scaffolded local apps set an app name too. + vi.stubEnv("DATABRICKS_APP_NAME", "local-app"); + const { req, next, status } = await invoke(); + expect(loadDevOboIdentityFromEnvironment).toHaveBeenCalledOnce(); + expect(req.headers).toMatchObject({ + "x-forwarded-access-token": "fake-token", + "x-forwarded-user": "alice", + "x-forwarded-email": "alice@example.com", + }); + expect(next).toHaveBeenCalledOnce(); + expect(status).not.toHaveBeenCalled(); + }); + + test.each(["production", "test", undefined])("is disabled in %s", (env) => { + vi.stubEnv("NODE_ENV", env); + expect(createDevOboMiddleware()).toBeUndefined(); + expect(loadDevOboIdentityFromEnvironment).not.toHaveBeenCalled(); + }); + + test.each([undefined, "", " "])( + "does not select a default profile when configured as %s", + (profile) => { + vi.stubEnv("DATABRICKS_CONFIG_PROFILE", profile); + expect(createDevOboMiddleware()).toBeUndefined(); + expect(loadDevOboIdentityFromEnvironment).not.toHaveBeenCalled(); + }, + ); + + test("enables injection with a token and host but no profile", async () => { + vi.stubEnv("DATABRICKS_CONFIG_PROFILE", undefined); + vi.stubEnv("DATABRICKS_HOST", "https://workspace.example"); + vi.stubEnv("DATABRICKS_TOKEN", "fake-direct-token"); + const { next } = await invoke(); + expect(next).toHaveBeenCalledOnce(); + expect(loadDevOboIdentityFromEnvironment).toHaveBeenCalledOnce(); + }); + + test("can be disabled to use the existing development fallback", () => { + vi.stubEnv("APPKIT_DEV_OBO", "false"); + expect(createDevOboMiddleware()).toBeUndefined(); + expect(loadDevOboIdentityFromEnvironment).not.toHaveBeenCalled(); + }); + + test.each(["localhost:8000", "127.0.0.1:8000", "[::1]:8000"])( + "accepts a same-origin request to %s", + async (host) => { + const req = makeRequest(); + req.headers = { + host, + origin: `http://${host}`, + "sec-fetch-site": "same-origin", + }; + const { next } = await invoke(req); + expect(next).toHaveBeenCalledOnce(); + }, + ); + + test.each([ + { host: "remote.example:8000" }, + { host: "localhost:9000" }, + { origin: "https://remote.example" }, + { origin: "http://localhost:9000" }, + { "sec-fetch-site": "cross-site" }, + { "sec-fetch-site": "same-site" }, + ])( + "does not load credentials for an untrusted request: %j", + async (headers) => { + const req = makeRequest(); + Object.assign(req.headers, headers); + const { status, next } = await invoke(req); + expect(status).toHaveBeenCalledWith(403); + expect(next).not.toHaveBeenCalled(); + expect(loadDevOboIdentityFromEnvironment).not.toHaveBeenCalled(); + }, + ); + + test("does not use proxy headers to decide whether a client is local", async () => { + const req = makeRequest(); + Object.defineProperty(req.socket, "remoteAddress", { value: "192.0.2.1" }); + req.headers["x-forwarded-for"] = "127.0.0.1"; + const { status, next } = await invoke(req); + expect(status).toHaveBeenCalledWith(403); + expect(next).not.toHaveBeenCalled(); + expect(loadDevOboIdentityFromEnvironment).not.toHaveBeenCalled(); + }); + + test("preserves explicitly forwarded credentials without loading the local profile", async () => { + const req = makeRequest(); + req.headers = { + host: "localhost:8000", + origin: "http://127.0.0.1:3001", + "x-forwarded-access-token": "forwarded-token", + "x-forwarded-user": "bob", + "x-forwarded-email": "bob@example.com", + }; + const expected = { ...req.headers }; + const { next } = await invoke(req); + expect(next).toHaveBeenCalledOnce(); + expect(req.headers).toEqual(expected); + expect(loadDevOboIdentityFromEnvironment).not.toHaveBeenCalled(); + }); + + test("clears an unrelated email when the resolved identity has none", async () => { + vi.mocked(loadDevOboIdentityFromEnvironment).mockResolvedValue({ + token: "fake-token", + userId: "alice", + }); + const req = makeRequest(); + req.headers["x-forwarded-email"] = "unrelated@example.com"; + await invoke(req); + expect(req.headers["x-forwarded-email"]).toBeUndefined(); + }); + + test("fails closed without exposing the CLI error or continuing as SP", async () => { + vi.mocked(loadDevOboIdentityFromEnvironment).mockRejectedValue( + new Error("credential-bearing failure"), + ); + const { req, status, json, next } = await invoke(); + expect(status).toHaveBeenCalledWith(401); + expect(json).toHaveBeenCalledWith({ + error: expect.stringContaining("DATABRICKS_TOKEN"), + }); + expect(JSON.stringify(json.mock.calls)).not.toContain("credential-bearing"); + expect(req.headers["x-forwarded-access-token"]).toBeUndefined(); + expect(next).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/appkit/src/plugins/server/tests/dev-obo.integration.test.ts b/packages/appkit/src/plugins/server/tests/dev-obo.integration.test.ts new file mode 100644 index 000000000..94b84da03 --- /dev/null +++ b/packages/appkit/src/plugins/server/tests/dev-obo.integration.test.ts @@ -0,0 +1,152 @@ +import type { Server } from "node:http"; + +import type { Router } from "express"; +import getPort from "get-port"; +import { loadDevOboIdentityFromEnvironment } from "shared"; +import { afterAll, beforeAll, describe, expect, test, vi } from "vitest"; + +import { getCurrentPrincipalKey } from "../../../context/execution-context"; +import { createRequestScope } from "../../../context/request-scope"; +import { AppKit, disposeApp } from "../../../core/appkit"; +import { Plugin, toPlugin } from "../../../plugin"; +import { defineManifest } from "../../../registry"; +import { getListeningPort, mockServiceContext } from "../../../testing"; +import { server } from "../index"; + +vi.mock("shared", async (importOriginal) => ({ + ...(await importOriginal()), + loadDevOboIdentityFromEnvironment: vi.fn(), +})); +vi.mock("../vite-dev-server", () => ({ + ViteDevServer: class { + async setup() {} + async close() {} + }, +})); + +class IdentityProbe extends Plugin { + static manifest = defineManifest<"identityProbe">({ + name: "identityProbe", + displayName: "Identity probe", + version: "1.0.0", + description: "Local execution identity test", + resources: { required: [], optional: [] }, + }); + + injectRoutes(router: Router) { + router.get("/user", (req, res) => { + createRequestScope(req).run(() => { + res.json({ principal: getCurrentPrincipalKey() }); + }); + }); + router.get("/app", (_req, res) => { + res.json({ principal: getCurrentPrincipalKey() }); + }); + } +} + +describe("normal dev server OBO injection", () => { + let baseUrl: string; + let httpServer: Server; + let dispose: () => Promise; + let service: ReturnType; + let now = Date.now(); + + beforeAll(async () => { + vi.stubEnv("NODE_ENV", "development"); + vi.stubEnv("DATABRICKS_CONFIG_PROFILE", "selected-user"); + vi.stubEnv("DATABRICKS_TOKEN", undefined); + vi.stubEnv("DATABRICKS_APP_NAME", "scaffolded-local-app"); + vi.stubEnv("APPKIT_DEV_OBO", undefined); + vi.spyOn(Date, "now").mockImplementation(() => now); + vi.mocked(loadDevOboIdentityFromEnvironment).mockResolvedValue({ + token: "fake-user-token", + userId: "alice", + email: "alice@example.com", + }); + service = mockServiceContext(); + const kit = await AppKit._createApp({ + plugins: [ + server({ + port: await getPort({ host: "127.0.0.1" }), + host: "127.0.0.1", + }), + toPlugin(IdentityProbe)({}), + ], + installSignalHandlers: false, + disableInternalTelemetry: true, + onPluginsReady(appkit) { + appkit.server.extend((app) => { + app.get("/custom-user", (req, res, next) => { + void appkit + .asUser(req) + .run(() => { + res.json({ principal: getCurrentPrincipalKey() }); + }) + .catch(next); + }); + }); + }, + }); + if (!(kit instanceof AppKit)) throw new Error("Expected AppKit"); + dispose = () => kit[disposeApp](); + httpServer = kit.server.getServer(); + baseUrl = `http://127.0.0.1:${await getListeningPort(httpServer)}`; + }); + + afterAll(async () => { + await dispose?.(); + service?.restore(); + vi.restoreAllMocks(); + vi.unstubAllEnvs(); + }); + + test("plugin routes run as the user without a proxy or manually supplied headers", async () => { + const response = await fetch(`${baseUrl}/api/identity-probe/user`); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ principal: "user:alice" }); + expect(service.createUserContextSpy).toHaveBeenCalledWith( + "fake-user-token", + "alice", + undefined, + "alice@example.com", + ); + expect(loadDevOboIdentityFromEnvironment).toHaveBeenCalledOnce(); + }); + + test("custom appkit.asUser routes get the same automatic headers", async () => { + const response = await fetch(`${baseUrl}/custom-user`); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ principal: "user:alice" }); + }); + + test("unscoped routes remain app-scoped despite the injected user headers", async () => { + const response = await fetch(`${baseUrl}/api/identity-probe/app`); + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ principal: "app" }); + }); + + test("refresh failure rejects and recovery never falls back to SP", async () => { + now += 31_000; + vi.mocked(loadDevOboIdentityFromEnvironment).mockRejectedValue( + new Error("credential-bearing failure"), + ); + const failed = await fetch(`${baseUrl}/api/identity-probe/user`); + expect(failed.status).toBe(401); + expect(await failed.text()).not.toContain("credential-bearing"); + + vi.mocked(loadDevOboIdentityFromEnvironment).mockResolvedValue({ + token: "refreshed-token", + userId: "alice", + }); + const recovered = await fetch(`${baseUrl}/api/identity-probe/user`); + expect(recovered.status).toBe(200); + expect(await recovered.json()).toEqual({ principal: "user:alice" }); + expect(service.createUserContextSpy).toHaveBeenLastCalledWith( + "refreshed-token", + "alice", + undefined, + undefined, + ); + }); +}); diff --git a/packages/appkit/src/stream/stream-manager.ts b/packages/appkit/src/stream/stream-manager.ts index 9d4840141..eade24ec2 100644 --- a/packages/appkit/src/stream/stream-manager.ts +++ b/packages/appkit/src/stream/stream-manager.ts @@ -3,8 +3,10 @@ import { randomUUID } from "node:crypto"; import { context } from "@opentelemetry/api"; import type { IAppResponse, StreamConfig } from "shared"; +import { normalizeIdentityError } from "../context/execution-context"; import { AppKitError } from "../errors/base"; import { ExecutionError } from "../errors/execution"; +import { IdentityExpiredError } from "../errors/identity-expired"; import { createLogger } from "../logging/logger"; import { EventRingBuffer } from "./buffers"; import { streamDefaults } from "./defaults"; @@ -303,7 +305,8 @@ export class StreamManager { } this._finalizeStream(streamEntry); - } catch (error) { + } catch (caught) { + const error = normalizeIdentityError(caught); // Two distinct messages: a *raw* one for server-side logs (full // detail, statement fragments, correlation IDs) and a *client* // one for the SSE payload (sanitized, stable, safe to render in @@ -318,7 +321,11 @@ export class StreamManager { // Upstream structured code (e.g. RESULT_TOO_LARGE_FOR_JSON_FALLBACK, // NOT_IMPLEMENTED). UI should branch on this, not on `error`. const upstreamCode = - error instanceof ExecutionError ? error.errorCode : undefined; + error instanceof IdentityExpiredError + ? error.code + : error instanceof ExecutionError + ? error.errorCode + : undefined; const errorEventId = randomUUID(); const errorCode = this._categorizeError(error); diff --git a/packages/appkit/src/telemetry/execution-span-processor.ts b/packages/appkit/src/telemetry/execution-span-processor.ts new file mode 100644 index 000000000..7441a0040 --- /dev/null +++ b/packages/appkit/src/telemetry/execution-span-processor.ts @@ -0,0 +1,32 @@ +import type { Context } from "@opentelemetry/api"; +import type { + ReadableSpan, + Span, + SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; + +import { + getCallerContext, + getCurrentActorId, +} from "../context/execution-context"; +import { ServiceContext } from "../context/service-context"; + +/** Attach identity at span creation, including cache, tool, and connector spans. */ +export class ExecutionSpanProcessor implements SpanProcessor { + onStart(span: Span, _parent: Context): void { + const caller = getCallerContext(); + span.setAttribute("appkit.execution.principal", caller ? "user" : "app"); + span.setAttribute( + "appkit.execution.principal_id", + caller?.principal.userId ?? + (ServiceContext.isInitialized() + ? ServiceContext.get().serviceUserId + : "app"), + ); + const actor = getCurrentActorId(); + if (actor) span.setAttribute("appkit.execution.actor_id", actor); + } + onEnd(_span: ReadableSpan): void {} + async forceFlush(): Promise {} + async shutdown(): Promise {} +} diff --git a/packages/appkit/src/telemetry/telemetry-manager.ts b/packages/appkit/src/telemetry/telemetry-manager.ts index a2b3c9d5f..c1fb43fc0 100644 --- a/packages/appkit/src/telemetry/telemetry-manager.ts +++ b/packages/appkit/src/telemetry/telemetry-manager.ts @@ -36,6 +36,7 @@ import { import type { TelemetryOptions } from "shared"; import { createLogger } from "../logging/logger"; +import { ExecutionSpanProcessor } from "./execution-span-processor"; import { TelemetryProvider } from "./telemetry-provider"; import { AppKitSampler } from "./trace-sampler"; import type { TelemetryConfig } from "./types"; @@ -205,7 +206,7 @@ export class TelemetryManager { this.tracerProvider = new NodeTracerProvider({ resource: this.resource, sampler: new AppKitSampler(), - spanProcessors: this.spanProcessors, + spanProcessors: [new ExecutionSpanProcessor(), ...this.spanProcessors], }); this.tracerProvider.register(); logger.debug( diff --git a/packages/appkit/src/telemetry/tests/execution-span-processor.test.ts b/packages/appkit/src/telemetry/tests/execution-span-processor.test.ts new file mode 100644 index 000000000..2cdacd868 --- /dev/null +++ b/packages/appkit/src/telemetry/tests/execution-span-processor.test.ts @@ -0,0 +1,62 @@ +import { ROOT_CONTEXT } from "@opentelemetry/api"; +import { + BasicTracerProvider, + InMemorySpanExporter, + SimpleSpanProcessor, +} from "@opentelemetry/sdk-trace-base"; +import { expect, test } from "vitest"; + +import { runInCallerContext } from "../../context"; +import { createMockWorkspaceClient } from "../../testing"; +import { ExecutionSpanProcessor } from "../execution-span-processor"; + +test("records app and caller identity on every span without tokens", async () => { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ + spanProcessors: [ + new ExecutionSpanProcessor(), + new SimpleSpanProcessor(exporter), + ], + }); + const tracer = provider.getTracer("identity"); + tracer.startSpan("app", {}, ROOT_CONTEXT).end(); + await Promise.all( + ["alice", "bob"].map((userId) => + runInCallerContext( + { + principal: { type: "user", userId }, + client: createMockWorkspaceClient(), + workspaceId: Promise.resolve("workspace"), + tokenFingerprint: "fingerprint-not-for-spans", + }, + async () => { + await Promise.resolve(); + for (const name of [ + "plugin.execute", + "tool.execute", + "cache.getOrExecute", + "connector.request", + ]) + tracer.startSpan(`${userId}:${name}`).end(); + }, + ), + ), + ); + const spans = exporter.getFinishedSpans(); + expect(spans[0].attributes).toEqual({ + "appkit.execution.principal": "app", + "appkit.execution.principal_id": "app", + }); + for (const span of spans.slice(1)) { + const userId = span.name.split(":")[0]; + expect(span.attributes).toEqual({ + "appkit.execution.principal": "user", + "appkit.execution.principal_id": userId, + "appkit.execution.actor_id": userId, + }); + } + expect(JSON.stringify(spans.map((s) => s.attributes))).not.toContain( + "fingerprint", + ); + await provider.shutdown(); +}); diff --git a/packages/shared/src/dev-obo.test.ts b/packages/shared/src/dev-obo.test.ts new file mode 100644 index 000000000..d2478a317 --- /dev/null +++ b/packages/shared/src/dev-obo.test.ts @@ -0,0 +1,173 @@ +import { afterEach, describe, expect, test, vi } from "vitest"; + +import { + createDevOboIdentityProvider, + loadDevOboIdentity, + loadDevOboIdentityFromEnvironment, +} from "./dev-obo"; + +afterEach(() => vi.restoreAllMocks()); + +describe("local OBO identity", () => { + test("passes the explicit profile to credential commands without leaking failures", async () => { + const run = vi.fn(async (args: string[]) => ({ + stdout: JSON.stringify( + args[0] === "auth" + ? { access_token: "fake-token" } + : { id: "alice", userName: "alice@example.com" }, + ), + })); + expect(await loadDevOboIdentity("chosen", run)).toEqual({ + token: "fake-token", + userId: "alice", + email: "alice@example.com", + }); + expect(run).toHaveBeenCalledWith(["auth", "token", "--profile", "chosen"]); + expect(run).toHaveBeenCalledWith([ + "current-user", + "me", + "--profile", + "chosen", + "--output", + "json", + ]); + run.mockRejectedValue(new Error("fake-token")); + const error = await loadDevOboIdentity("chosen", run).catch((err) => err); + expect(String(error)).not.toContain("fake-token"); + expect(error.cause).toBeUndefined(); + }); + + test("prefers DATABRICKS_TOKEN and resolves its user without a profile", async () => { + const run = vi.fn(async (_args: string[], _env?: NodeJS.ProcessEnv) => ({ + stdout: JSON.stringify({ + id: "alice", + userName: "alice@example.com", + }), + })); + + await expect( + loadDevOboIdentityFromEnvironment( + { + DATABRICKS_HOST: "https://workspace.example", + DATABRICKS_TOKEN: "fake-direct-token", + DATABRICKS_CONFIG_PROFILE: "ignored-profile", + }, + run, + ), + ).resolves.toEqual({ + token: "fake-direct-token", + userId: "alice", + email: "alice@example.com", + }); + expect(run).toHaveBeenCalledOnce(); + expect(run.mock.calls[0]?.[0]).toEqual([ + "current-user", + "me", + "--host", + "https://workspace.example", + "--output", + "json", + ]); + expect(run.mock.calls[0]?.[1]).toMatchObject({ + DATABRICKS_HOST: "https://workspace.example", + DATABRICKS_TOKEN: "fake-direct-token", + }); + expect(run.mock.calls[0]?.[1]).not.toHaveProperty( + "DATABRICKS_CONFIG_PROFILE", + ); + }); + + test("falls back to the explicitly configured profile without a token", async () => { + const run = vi.fn(async (args: string[]) => ({ + stdout: JSON.stringify( + args[0] === "auth" + ? { access_token: "fake-profile-token" } + : { id: "alice" }, + ), + })); + + await expect( + loadDevOboIdentityFromEnvironment( + { DATABRICKS_CONFIG_PROFILE: "selected-user" }, + run, + ), + ).resolves.toEqual({ token: "fake-profile-token", userId: "alice" }); + expect(run).toHaveBeenCalledWith([ + "auth", + "token", + "--profile", + "selected-user", + ]); + }); + + test("does not fall back to a profile when a token has no host", async () => { + const run = vi.fn(); + await expect( + loadDevOboIdentityFromEnvironment( + { + DATABRICKS_TOKEN: "fake-direct-token", + DATABRICKS_CONFIG_PROFILE: "must-not-be-used", + }, + run, + ), + ).rejects.toThrow("Set DATABRICKS_HOST"); + expect(run).not.toHaveBeenCalled(); + }); + + test("never falls back to an implicit CLI profile", async () => { + const run = vi.fn(); + await expect(loadDevOboIdentity(" ", run)).rejects.toThrow( + "explicitly selected Databricks user profile", + ); + expect(run).not.toHaveBeenCalled(); + }); + + test.each([ + { id: "sp", applicationId: "application-id" }, + { + id: "sp", + schemas: ["urn:ietf:params:scim:schemas:core:2.0:ServicePrincipal"], + }, + ])("rejects a service principal profile: %j", async (user) => { + const run = vi.fn(async (args: string[]) => ({ + stdout: JSON.stringify( + args[0] === "auth" ? { access_token: "fake-token" } : user, + ), + })); + await expect(loadDevOboIdentity("selected-sp", run)).rejects.toThrow( + "user profile", + ); + }); + + test("shares concurrent initial loads and refreshes without using stale credentials", async () => { + let now = 1_000; + vi.spyOn(Date, "now").mockImplementation(() => now); + const first = { token: "fake-token-1", userId: "alice" }; + const second = { token: "fake-token-2", userId: "alice" }; + const load = vi.fn().mockResolvedValue(first); + const getIdentity = createDevOboIdentityProvider(load); + expect(await Promise.all([getIdentity(), getIdentity()])).toEqual([ + first, + first, + ]); + expect(load).toHaveBeenCalledTimes(1); + + now += 29_999; + expect(await getIdentity()).toEqual(first); + expect(load).toHaveBeenCalledTimes(1); + now += 1; + load.mockResolvedValue(second); + expect(await Promise.all([getIdentity(), getIdentity()])).toEqual([ + second, + second, + ]); + expect(load).toHaveBeenCalledTimes(2); + + now += 30_000; + load.mockRejectedValue(new Error("expired")); + await expect(getIdentity()).rejects.toThrow("expired"); + await expect(getIdentity()).rejects.toThrow("expired"); + load.mockResolvedValue(second); + expect(await getIdentity()).toEqual(second); + }); +}); diff --git a/packages/shared/src/dev-obo.ts b/packages/shared/src/dev-obo.ts new file mode 100644 index 000000000..80d96899b --- /dev/null +++ b/packages/shared/src/dev-obo.ts @@ -0,0 +1,107 @@ +import { execFile } from "node:child_process"; +import { promisify } from "node:util"; + +const execute = promisify(execFile); + +type DatabricksCliRunner = ( + args: string[], + env?: NodeJS.ProcessEnv, +) => Promise<{ stdout: string }>; + +const runDatabricksCli: DatabricksCliRunner = (args, env) => + execute("databricks", args, { env, timeout: 30_000 }); + +export interface DevOboIdentity { + token: string; + userId: string; + email?: string; +} + +function identityFromUser(token: string, userOutput: string): DevOboIdentity { + const user = JSON.parse(userOutput); + if ( + !token.trim() || + typeof user.id !== "string" || + !user.id.trim() || + user.applicationId || + user.schemas?.some((schema: string) => schema.endsWith(":ServicePrincipal")) + ) + throw new Error(); + return { + token, + userId: user.id, + email: typeof user.userName === "string" ? user.userName : undefined, + }; +} + +/** CLI output is credential-bearing. Never include it in an error or log. */ +export async function loadDevOboIdentity( + profile: string, + run: DatabricksCliRunner = runDatabricksCli, +): Promise { + try { + if (!profile.trim()) throw new Error(); + const [tokenOutput, userOutput] = await Promise.all([ + run(["auth", "token", "--profile", profile]), + run(["current-user", "me", "--profile", profile, "--output", "json"]), + ]); + const token = JSON.parse(tokenOutput.stdout).access_token; + if (typeof token !== "string") throw new Error(); + return identityFromUser(token, userOutput.stdout); + } catch { + throw new Error( + "Unable to obtain local OBO credentials. Authenticate the explicitly selected Databricks user profile and retry.", + ); + } +} + +/** Prefer explicit token credentials, then fall back to an explicit profile. */ +export async function loadDevOboIdentityFromEnvironment( + env: NodeJS.ProcessEnv = process.env, + run: DatabricksCliRunner = runDatabricksCli, +): Promise { + const token = env.DATABRICKS_TOKEN?.trim(); + if (!token) { + return loadDevOboIdentity(env.DATABRICKS_CONFIG_PROFILE ?? "", run); + } + + try { + const host = env.DATABRICKS_HOST?.trim(); + if (!host) throw new Error(); + const tokenEnvironment = { ...env }; + delete tokenEnvironment.DATABRICKS_CONFIG_PROFILE; + tokenEnvironment.DATABRICKS_HOST = host; + tokenEnvironment.DATABRICKS_TOKEN = token; + const userOutput = await run( + ["current-user", "me", "--host", host, "--output", "json"], + tokenEnvironment, + ); + return identityFromUser(token, userOutput.stdout); + } catch { + throw new Error( + "Unable to obtain local OBO identity from DATABRICKS_TOKEN. Set DATABRICKS_HOST for the same workspace and verify the token belongs to a user.", + ); + } +} + +/** Share concurrent refreshes and keep credentials only in memory. */ +export function createDevOboIdentityProvider( + load: () => Promise, +): () => Promise { + let identity: DevOboIdentity | undefined; + let loadedAt = 0; + let refresh: Promise | undefined; + return async () => { + if (identity && Date.now() - loadedAt < 30_000) return identity; + refresh ??= load() + .then((next) => { + identity = next; + loadedAt = Date.now(); + return next; + }) + .finally(() => { + refresh = undefined; + }); + return refresh; + }; +} diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index c8e7e8fa5..42abeb903 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -1,5 +1,10 @@ export * from "./agent"; export * from "./cache"; +export { + createDevOboIdentityProvider, + loadDevOboIdentity, + loadDevOboIdentityFromEnvironment, +} from "./dev-obo"; export * from "./execute"; export * from "./genie"; export * from "./metric-filter";