@databricks/appkit 0.83.0 → 0.84.0

package/dist/plugins/agents/agents.js

@@ -1,6 +1,10 @@
+import { AuthenticationError } from "../../errors/authentication.js";
+import { IdentityExpiredError } from "../../errors/identity-expired.js";
+import "../../errors/index.js";
 import { createLogger } from "../../logging/logger.js";
-import { getWorkspaceClient } from "../../context/execution-context.js";
+import { getWorkspaceClient, normalizeIdentityError, runOutsideCallerScope } from "../../context/execution-context.js";
 import "../../context/index.js";
+import { createRequestScope } from "../../context/request-scope.js";
 import { Plugin } from "../../plugin/plugin.js";
 import { toPlugin } from "../../plugin/to-plugin.js";
 import "../../plugin/index.js";
@@ -22,6 +26,7 @@ import { loadAgentsFromDir } from "../../core/agent/load-agents.js";
 import { ActiveStreamTracker } from "./active-stream-tracker.js";
 import { buildAdapterExtensions, supervisorToolDescription, warnOnCapabilityMismatch } from "./adapter-extensions.js";
 import { requiresApproval } from "./approval.js";
+import { isOnBehalfOfUser, modelClientProvider, requireOboCaller, runInOboAgentRun } from "./auth-mode.js";
 import { LOAD_SKILL_TOOL_DEF, READ_SKILL_FILE_TOOL_DEF } from "./builtin-tools.js";
 import { agentStreamDefaults } from "./defaults.js";
 import { EventChannel } from "./event-channel.js";
@@ -276,6 +281,8 @@ var AgentsPlugin = class extends Plugin {
 	}
 	async buildRegisteredAgent(name, def, src) {
 		const adapter = await this.resolveAdapter(def, name);
+		const auth = def.auth ?? this.config.auth;
+		if (isOnBehalfOfUser(auth) && (await import("../../agents/databricks.js")).hasFixedWorkspaceClient(adapter)) throw new Error(`Agent '${name}' is on-behalf-of-user, but its model adapter has a fixed workspaceClient, so the model would run as the service principal. Pass \`workspaceClient: () => getWorkspaceClient()\` to the adapter, or use a model string.`);
 		const skills = await this.resolveAgentSkills(name, def, src);
 		const toolIndex = await this.buildToolIndex(name, def, src, skills);
 		warnOnCapabilityMismatch(name, adapter, toolIndex);
@@ -289,7 +296,8 @@ var AgentsPlugin = class extends Plugin {
 			maxTokens: def.maxTokens,
 			generationParams: def.generationParams,
 			ephemeral: def.ephemeral,
-			skills
+			skills,
+			...auth && { auth }
 		};
 	}
 	/** Loads the shared global skill pool from `<agentsDir>/skills/`. */
@@ -329,7 +337,10 @@ var AgentsPlugin = class extends Plugin {
 		if (!source) throw new Error(`Agent '${name}' has no model configured and no DATABRICKS_SERVING_ENDPOINT_NAME default available`);
 		if (typeof source === "string") {
 			const { adapterFromModelString } = await import("../../agents/databricks.js");
-			return adapterFromModelString(source, adapterOptions);
+			return adapterFromModelString(source, {
+				...adapterOptions,
+				workspaceClient: modelClientProvider()
+			});
 		}
 		return await source;
 	}
@@ -629,6 +640,7 @@ var AgentsPlugin = class extends Plugin {
 			res.status(400).json({ error: agentName ? `Agent "${agentName}" not found` : "No agent registered" });
 			return;
 		}
+		if (!this.allowsCaller(registered, req, res)) return;
 		const userId = this.resolveUserId(req);
 		const limits = this.resolvedLimits;
 		if (this.countUserStreams(userId) >= limits.maxConcurrentStreamsPerUser) {
@@ -656,7 +668,7 @@ var AgentsPlugin = class extends Plugin {
 			res.status(500).json({ error: "Thread operation failed" });
 			return;
 		}
-		return this._streamAgent(req, res, registered, thread, userId, mlflowRunId, skill);
+		return this.runInAgentScope(registered, req, () => this._streamAgent(req, res, registered, thread, userId, mlflowRunId, skill));
 	}
 	/**
 	* Returns the names of tools in `registered.toolIndex` whose annotations
@@ -700,6 +712,7 @@ var AgentsPlugin = class extends Plugin {
 			res.status(400).json({ error: "No agent registered" });
 			return;
 		}
+		if (!this.allowsCaller(registered, req, res)) return;
 		const approvalGated = this.collectApprovalRequiredToolNames(registered);
 		if (approvalGated.length > 0) {
 			res.status(400).json({ error: `Agent '${registered.name}' exposes ${approvalGated.length} approval-gated tool(s) (${approvalGated.join(", ")}); /invocations and /responses are non-streaming and cannot run HITL. Use POST /chat for HITL-capable agents, or disable approval via agents({ approval: { requireForDestructive: false } }).` });
@@ -737,7 +750,35 @@ var AgentsPlugin = class extends Plugin {
 			res.status(500).json({ error: "Thread operation failed" });
 			return;
 		}
-		return this._runAgentNonStreaming(req, res, registered, thread, userId, mlflowRunId);
+		return this.runInAgentScope(registered, req, () => this._runAgentNonStreaming(req, res, registered, thread, userId, mlflowRunId));
+	}
+	/**
+	* Fail closed for an on-behalf-of-user agent: without a forwarded user token
+	* the request is rejected with 401 before any model or tool call, in
+	* production and in development. Other agents are unaffected.
+	*/
+	allowsCaller(registered, req, res) {
+		if (!isOnBehalfOfUser(registered.auth)) return true;
+		try {
+			requireOboCaller(req);
+			return true;
+		} catch (error) {
+			const failure = error instanceof AuthenticationError ? error : AuthenticationError.missingToken("user token");
+			res.status(401).json({
+				error: failure.clientMessage,
+				code: failure.code
+			});
+			return false;
+		}
+	}
+	/**
+	* Run an on-behalf-of-user agent inside the request's user scope, so the
+	* model call, hand-rolled tools, and sub-agents all act as the user. Other
+	* agents run exactly as before.
+	*/
+	runInAgentScope(registered, req, fn) {
+		if (!isOnBehalfOfUser(registered.auth)) return fn();
+		return createRequestScope(req).run(() => runInOboAgentRun(fn));
 	}
 	async _streamAgent(req, res, registered, thread, userId, mlflowRunId, forcedSkill) {
 		const abortController = new AbortController();
@@ -811,12 +852,12 @@ var AgentsPlugin = class extends Plugin {
 							role: "assistant",
 							content: fullContent
 						});
-						await this.threadStore.addMessage(thread.id, userId, {
+						await runOutsideCallerScope(() => this.threadStore.addMessage(thread.id, userId, {
 							id: randomUUID(),
 							role: "assistant",
 							content: fullContent,
 							createdAt: /* @__PURE__ */ new Date()
-						});
+						}));
 					}
 					const mlflowTraceId = currentTraceId();
 					if (mlflowTraceId) for (const evt of translator.translate({
@@ -830,14 +871,15 @@ var AgentsPlugin = class extends Plugin {
 					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 {
 				this.approvalGate.abortStream(requestId);
 				this.untrackStream(requestId);
 				if (registered.ephemeral) try {
-					await this.threadStore.delete(thread.id, userId);
+					await runOutsideCallerScope(() => this.threadStore.delete(thread.id, userId));
 				} catch (err) {
 					logger.warn("Failed to delete ephemeral thread %s: %O", thread.id, err);
 				}
@@ -939,12 +981,12 @@ var AgentsPlugin = class extends Plugin {
 						role: "assistant",
 						content: fullContent
 					});
-					await this.threadStore.addMessage(thread.id, userId, {
+					await runOutsideCallerScope(() => this.threadStore.addMessage(thread.id, userId, {
 						id: randomUUID(),
 						role: "assistant",
 						content: fullContent,
 						createdAt: /* @__PURE__ */ new Date()
-					});
+					}));
 				}
 				mlflowTraceId = currentTraceId();
 			});
@@ -953,7 +995,15 @@ var AgentsPlugin = class extends Plugin {
 				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" : error instanceof Error ? error.message : String(error);
 			res.status(500).json({ error: message });
 			return;
@@ -961,7 +1011,7 @@ var AgentsPlugin = class extends Plugin {
 			this.approvalGate.abortStream(requestId);
 			this.untrackStream(requestId);
 			if (registered.ephemeral) try {
-				await this.threadStore.delete(thread.id, userId);
+				await runOutsideCallerScope(() => this.threadStore.delete(thread.id, userId));
 			} catch (err) {
 				logger.warn("Failed to delete ephemeral thread %s: %O", thread.id, err);
 			}
@@ -989,7 +1039,7 @@ var AgentsPlugin = class extends Plugin {
 		});
 	}
 	dispatchSkillTool(entry, args) {
-		return dispatchSkillTool(entry, args, () => this.skillWorkspaceClient());
+		return runOutsideCallerScope(() => dispatchSkillTool(entry, args, () => this.skillWorkspaceClient()));
 	}
 	/**
 	* Renders the prompt addendum for a force-loaded skill (`/skill-name`).