@databricks/appkit 0.32.0 → 0.33.0

package/dist/plugins/agents/agents.js

@@ -6,6 +6,8 @@ import { buildMcpHostPolicy } from "../../connectors/mcp/host-policy.js";
 import { AppKitMcpClient } from "../../connectors/mcp/client.js";
 import "../../connectors/mcp/index.js";
 import { consumeAdapterStream } from "../../core/agent/consume-adapter-stream.js";
+import { createPluginsProxy } from "../../core/agent/plugins-map.js";
+import { resolveToolkitFromProvider } from "../../core/agent/toolkit-resolver.js";
 import { functionToolToDefinition, isFunctionTool } from "../../core/agent/tools/function-tool.js";
 import { isHostedTool, resolveHostedTools } from "../../core/agent/tools/hosted-tools.js";
 import { isToolkitEntry } from "../../core/agent/types.js";
@@ -56,6 +58,13 @@ var AgentsPlugin = class extends Plugin {
 	agents = /* @__PURE__ */ new Map();
 	defaultAgentName = null;
 	activeStreams = /* @__PURE__ */ new Map();
+	/**
+	* Per-user stream count, kept in sync with `activeStreams` so the
+	* concurrent-stream rate limit check is O(1) instead of O(n) over every
+	* active stream on every request. Mutated only via {@link trackStream}
+	* and {@link untrackStream}.
+	*/
+	userStreamCounts = /* @__PURE__ */ new Map();
 	mcpClient = null;
 	threadStore;
 	approvalGate = new ToolApprovalGate();
@@ -69,13 +78,32 @@ var AgentsPlugin = class extends Plugin {
 			else logger.info("Using default InMemoryThreadStore (dev-only — threads are lost on restart and grow without bound).");
 		}
 	}
-	/** Effective approval policy with defaults applied. */
+	/**
+	* Effective approval policy with defaults applied. Memoised so the
+	* `timeoutMs` validation warning fires at most once per plugin instance —
+	* `resolvedApprovalPolicy` gets hit on every chat stream and a noisy
+	* misconfig would otherwise spam the logs.
+	*
+	* `timeoutMs` is clamped to a 1s floor so a misconfigured value (`0`,
+	* negative, or `NaN`) can't degrade into immediate auto-denial of every
+	* mutating tool call.
+	*/
+	cachedApprovalPolicy = null;
 	get resolvedApprovalPolicy() {
+		if (this.cachedApprovalPolicy) return this.cachedApprovalPolicy;
 		const cfg = this.config.approval ?? {};
-		return {
+		const APPROVAL_TIMEOUT_FLOOR_MS = 1e3;
+		const APPROVAL_TIMEOUT_DEFAULT_MS = 6e4;
+		let timeoutMs = cfg.timeoutMs ?? APPROVAL_TIMEOUT_DEFAULT_MS;
+		if (!Number.isFinite(timeoutMs) || timeoutMs < APPROVAL_TIMEOUT_FLOOR_MS) {
+			logger.warn("approval.timeoutMs=%s is below the %sms floor; using default %sms instead. Mutating tool calls would otherwise auto-deny before any UI could respond.", cfg.timeoutMs, APPROVAL_TIMEOUT_FLOOR_MS, APPROVAL_TIMEOUT_DEFAULT_MS);
+			timeoutMs = APPROVAL_TIMEOUT_DEFAULT_MS;
+		}
+		this.cachedApprovalPolicy = {
 			requireForDestructive: cfg.requireForDestructive ?? true,
-			timeoutMs: cfg.timeoutMs ?? 6e4
+			timeoutMs
 		};
+		return this.cachedApprovalPolicy;
 	}
 	/** Effective DoS limits with defaults applied. */
 	get resolvedLimits() {
@@ -87,11 +115,35 @@ var AgentsPlugin = class extends Plugin {
 			toolCallTimeoutMs: cfg.toolCallTimeoutMs ?? 3e5
 		};
 	}
-	/** Count active streams owned by a given user. */
+	/** Count active streams owned by a given user. O(1). */
 	countUserStreams(userId) {
-		let n = 0;
-		for (const entry of this.activeStreams.values()) if (entry.userId === userId) n++;
-		return n;
+		return this.userStreamCounts.get(userId) ?? 0;
+	}
+	/**
+	* Register a stream for `userId` and bump the per-user counter. Paired
+	* with {@link untrackStream}; the two helpers are the only writers to
+	* `activeStreams` + `userStreamCounts`, so the counter cannot drift from
+	* the map.
+	*/
+	trackStream(requestId, userId, controller) {
+		this.activeStreams.set(requestId, {
+			controller,
+			userId
+		});
+		this.userStreamCounts.set(userId, (this.userStreamCounts.get(userId) ?? 0) + 1);
+	}
+	/**
+	* Remove a stream from the active map and decrement the per-user
+	* counter. Idempotent — calling twice for the same `requestId` is a
+	* no-op (the second call sees no entry and returns early).
+	*/
+	untrackStream(requestId) {
+		const entry = this.activeStreams.get(requestId);
+		if (!entry) return;
+		this.activeStreams.delete(requestId);
+		const next = (this.userStreamCounts.get(entry.userId) ?? 0) - 1;
+		if (next <= 0) this.userStreamCounts.delete(entry.userId);
+		else this.userStreamCounts.set(entry.userId, next);
 	}
 	async setup() {
 		const { agents, defaultAgentName } = await this.buildAgentRegistry();
@@ -108,10 +160,6 @@ var AgentsPlugin = class extends Plugin {
 	*/
 	async reload() {
 		const next = await this.buildAgentRegistry();
-		if (this.mcpClient) {
-			await this.mcpClient.close();
-			this.mcpClient = null;
-		}
 		this.agents = next.agents;
 		this.defaultAgentName = next.defaultAgentName;
 	}
@@ -180,7 +228,9 @@ var AgentsPlugin = class extends Plugin {
 	}
 	/**
 	* Builds the map of plugin-name → toolkit that the markdown loader consults
-	* when resolving `toolkits:` frontmatter entries.
+	* when resolving `plugin:NAME` entries in the unified `tools:` frontmatter
+	* list (and, equivalently, that the code form passes as the `plugins`
+	* argument to `tools(plugins) => Record<...>`).
 	*/
 	pluginProviderIndex() {
 		const out = /* @__PURE__ */ new Map();
@@ -215,7 +265,7 @@ var AgentsPlugin = class extends Plugin {
 			try {
 				return await DatabricksAdapter.fromModelServing(void 0, adapterOptions);
 			} catch (err) {
-				throw new Error(`Agent '${name}' has no model configured and no DATABRICKS_AGENT_ENDPOINT default available`, { cause: err instanceof Error ? err : void 0 });
+				throw new Error(`Agent '${name}' has no model configured and no DATABRICKS_SERVING_ENDPOINT_NAME default available`, { cause: err instanceof Error ? err : void 0 });
 			}
 		}
 		if (typeof source === "string") {
@@ -231,10 +281,11 @@ var AgentsPlugin = class extends Plugin {
 	*/
 	async buildToolIndex(agentName, def, src) {
 		const index = /* @__PURE__ */ new Map();
-		const hasExplicitTools = def.tools && Object.keys(def.tools).length > 0;
+		const hasDeclaredTools = def.tools !== void 0;
+		const toolsRecord = this.resolveDefTools(agentName, def);
 		const hasExplicitSubAgents = def.agents && Object.keys(def.agents).length > 0;
 		const inheritDefaults = normalizeAutoInherit(this.config.autoInheritTools);
-		if (!hasExplicitTools && !hasExplicitSubAgents && (src.origin === "file" ? inheritDefaults.file : inheritDefaults.code)) await this.applyAutoInherit(agentName, index);
+		if (!hasDeclaredTools && !hasExplicitSubAgents && (src.origin === "file" ? inheritDefaults.file : inheritDefaults.code)) await this.applyAutoInherit(agentName, index);
 		for (const [childKey, childDef] of Object.entries(def.agents ?? {})) {
 			const toolName = `agent-${childKey}`;
 			index.set(toolName, {
@@ -255,7 +306,7 @@ var AgentsPlugin = class extends Plugin {
 			});
 		}
 		const hostedToCollect = [];
-		for (const [key, tool] of Object.entries(def.tools ?? {})) {
+		for (const [key, tool] of Object.entries(toolsRecord)) {
 			if (isToolkitEntry(tool)) {
 				index.set(key, {
 					source: "toolkit",
@@ -288,6 +339,44 @@ var AgentsPlugin = class extends Plugin {
 		if (hostedToCollect.length > 0) await this.connectHostedTools(hostedToCollect, index);
 		return index;
 	}
+	/**
+	* Resolves an `AgentDefinition.tools` field to a plain tool record. The
+	* function form is invoked exactly once at agent setup with the typed
+	* {@link Plugins} map; the result replaces the function reference for the
+	* remainder of the registered agent's lifetime.
+	*
+	* Plain object form is returned as-is; an undefined `tools` returns an
+	* empty record. The function form is wrapped in a try/catch so a thrown
+	* callback fails registration with a useful message instead of leaking
+	* the raw stack.
+	*/
+	resolveDefTools(agentName, def) {
+		if (typeof def.tools !== "function") return def.tools ?? {};
+		try {
+			return def.tools(this.buildPluginsMap());
+		} catch (err) {
+			throw new Error(`Agent '${agentName}': tools(plugins) callback threw: ${err instanceof Error ? err.message : String(err)}`, { cause: err instanceof Error ? err : void 0 });
+		}
+	}
+	/**
+	* Builds the typed {@link Plugins} map passed to the function form of
+	* `AgentDefinition.tools`. Each entry exposes the plugin instance directly
+	* (so user code can call typed instance methods including `.toolkit()`);
+	* plugins missing `.toolkit()` get a synthesized fallback that walks
+	* `getAgentTools()` via `resolveToolkitFromProvider`.
+	*
+	* Wrapped in {@link createPluginsProxy} so that accessing an unknown
+	* plugin name throws a named "not registered, Available: ..." error
+	* instead of bubbling up a generic `Cannot read properties of undefined`
+	* from the agent's `tools(plugins)` callback.
+	*/
+	buildPluginsMap() {
+		const out = {};
+		if (!this.context) return createPluginsProxy(out, `Agent '${this.name}': tools(plugins)`);
+		for (const { name, provider } of this.context.getToolProviders()) if (typeof provider.toolkit === "function") out[name] = provider;
+		else out[name] = { toolkit: (opts) => resolveToolkitFromProvider(name, provider, opts) };
+		return createPluginsProxy(out, `Agent '${this.name}': tools(plugins)`);
+	}
 	async applyAutoInherit(agentName, index) {
 		if (!this.context) return;
 		const inherited = [];
@@ -299,29 +388,23 @@ var AgentsPlugin = class extends Plugin {
 		};
 		for (const { name: pluginName, provider } of this.context.getToolProviders()) {
 			if (pluginName === this.name) continue;
-			const withToolkit = provider;
-			if (typeof withToolkit.toolkit === "function") {
-				const entries = withToolkit.toolkit();
-				for (const [key, maybeEntry] of Object.entries(entries)) {
-					if (!isToolkitEntry(maybeEntry)) continue;
-					if (maybeEntry.autoInheritable !== true) {
-						recordSkip(maybeEntry.pluginName, maybeEntry.localName);
-						continue;
-					}
-					index.set(key, {
-						source: "toolkit",
-						pluginName: maybeEntry.pluginName,
-						localName: maybeEntry.localName,
-						def: {
-							...maybeEntry.def,
-							name: key
-						}
-					});
-					inherited.push(key);
+			const entries = resolveToolkitFromProvider(pluginName, provider);
+			for (const [key, entry] of Object.entries(entries)) {
+				if (entry.autoInheritable !== true) {
+					recordSkip(entry.pluginName, entry.localName);
+					continue;
 				}
-				continue;
+				index.set(key, {
+					source: "toolkit",
+					pluginName: entry.pluginName,
+					localName: entry.localName,
+					def: {
+						...entry.def,
+						name: key
+					}
+				});
+				inherited.push(key);
 			}
-			for (const tool of provider.getAgentTools()) recordSkip(pluginName, tool.name);
 		}
 		if (inherited.length > 0) logger.info("[agent %s] auto-inherited %d tool(s): %s", agentName, inherited.length, inherited.join(", "));
 		if (skippedByPlugin.size > 0) {
@@ -358,7 +441,8 @@ var AgentsPlugin = class extends Plugin {
 			this.mcpClient = new AppKitMcpClient(host, authenticate, policy);
 		}
 		const endpoints = resolveHostedTools(hostedTools);
-		await this.mcpClient.connectAll(endpoints);
+		const result = await this.mcpClient.connectAll(endpoints);
+		if (result.failed.length > 0) logger.warn("MCP: %s of %s endpoints failed to connect (%s). Agents that reference these endpoints will boot without their hosted tools.", result.failed.length, endpoints.length, result.failed.map((f) => f.name).join(", "));
 		for (const def of this.mcpClient.getAllToolDefinitions()) index.set(def.name, {
 			source: "mcp",
 			mcpToolName: def.name,
@@ -454,19 +538,26 @@ var AgentsPlugin = class extends Plugin {
 			res.status(429).json({ error: `Too many concurrent streams for this user (limit ${limits.maxConcurrentStreamsPerUser}). Wait for an existing stream to complete before starting another.` });
 			return;
 		}
-		let thread = threadId ? await this.threadStore.get(threadId, userId) : null;
-		if (threadId && !thread) {
-			res.status(404).json({ error: `Thread ${threadId} not found` });
+		let thread;
+		try {
+			const existing = threadId ? await this.threadStore.get(threadId, userId) : null;
+			if (threadId && !existing) {
+				res.status(404).json({ error: `Thread ${threadId} not found` });
+				return;
+			}
+			thread = existing ?? await this.threadStore.create(userId);
+			const userMessage = {
+				id: randomUUID(),
+				role: "user",
+				content: message,
+				createdAt: /* @__PURE__ */ new Date()
+			};
+			await this.threadStore.addMessage(thread.id, userId, userMessage);
+		} catch (err) {
+			logger.error("threadStore failed in /chat: %O", err);
+			res.status(500).json({ error: "Thread operation failed" });
 			return;
 		}
-		if (!thread) thread = await this.threadStore.create(userId);
-		const userMessage = {
-			id: randomUUID(),
-			role: "user",
-			content: message,
-			createdAt: /* @__PURE__ */ new Date()
-		};
-		await this.threadStore.addMessage(thread.id, userId, userMessage);
 		return this._streamAgent(req, res, registered, thread, userId);
 	}
 	async _handleInvocations(req, res) {
@@ -491,23 +582,30 @@ var AgentsPlugin = class extends Plugin {
 			res.status(429).json({ error: `Too many concurrent streams for this user (limit ${limits.maxConcurrentStreamsPerUser}). Wait for an existing stream to complete before starting another.` });
 			return;
 		}
-		const thread = await this.threadStore.create(userId);
-		if (typeof input === "string") await this.threadStore.addMessage(thread.id, userId, {
-			id: randomUUID(),
-			role: "user",
-			content: input,
-			createdAt: /* @__PURE__ */ new Date()
-		});
-		else for (const item of input) {
-			const role = item.role ?? "user";
-			const content = typeof item.content === "string" ? item.content : JSON.stringify(item.content ?? "");
-			if (!content) continue;
-			await this.threadStore.addMessage(thread.id, userId, {
+		let thread;
+		try {
+			thread = await this.threadStore.create(userId);
+			if (typeof input === "string") await this.threadStore.addMessage(thread.id, userId, {
 				id: randomUUID(),
-				role,
-				content,
+				role: "user",
+				content: input,
 				createdAt: /* @__PURE__ */ new Date()
 			});
+			else for (const item of input) {
+				const role = item.role ?? "user";
+				const content = typeof item.content === "string" ? item.content : JSON.stringify(item.content ?? "");
+				if (!content) continue;
+				await this.threadStore.addMessage(thread.id, userId, {
+					id: randomUUID(),
+					role,
+					content,
+					createdAt: /* @__PURE__ */ new Date()
+				});
+			}
+		} catch (err) {
+			logger.error("threadStore failed in /invocations: %O", err);
+			res.status(500).json({ error: "Thread operation failed" });
+			return;
 		}
 		return this._streamAgent(req, res, registered, thread, userId);
 	}
@@ -515,10 +613,7 @@ var AgentsPlugin = class extends Plugin {
 		const abortController = new AbortController();
 		const signal = abortController.signal;
 		const requestId = randomUUID();
-		this.activeStreams.set(requestId, {
-			controller: abortController,
-			userId
-		});
+		this.trackStream(requestId, userId, abortController);
 		const tools = Array.from(registered.toolIndex.values()).map((e) => e.def);
 		const approvalPolicy = this.resolvedApprovalPolicy;
 		const limits = this.resolvedLimits;
@@ -585,7 +680,7 @@ var AgentsPlugin = class extends Plugin {
 				return;
 			} finally {
 				this.approvalGate.abortStream(requestId);
-				this.activeStreams.delete(requestId);
+				this.untrackStream(requestId);
 				if (registered.ephemeral) try {
 					await this.threadStore.delete(thread.id, userId);
 				} catch (err) {
@@ -648,8 +743,10 @@ var AgentsPlugin = class extends Plugin {
 		if (entry.source === "toolkit") {
 			if (!this.context) throw new Error("Plugin tool execution requires PluginContext; this should never happen through createApp");
 			result = await this.context.executeTool(runState.req, entry.pluginName, entry.localName, args, runState.signal, runState.limits.toolCallTimeoutMs);
-		} else if (entry.source === "function") result = await entry.functionTool.execute(args);
-		else if (entry.source === "mcp") {
+		} else if (entry.source === "function") {
+			if (typeof args !== "object" || args === null || Array.isArray(args)) throw new Error(`Function tool '${name}' received non-object arguments (got ${args === null ? "null" : Array.isArray(args) ? "array" : typeof args}); expected a JSON object.`);
+			result = await entry.functionTool.execute(args);
+		} else if (entry.source === "mcp") {
 			if (!this.mcpClient) throw new Error("MCP client not connected");
 			const oboToken = runState.req.headers["x-forwarded-access-token"];
 			const mcpAuth = typeof oboToken === "string" ? { Authorization: `Bearer ${oboToken}` } : void 0;
@@ -734,7 +831,7 @@ var AgentsPlugin = class extends Plugin {
 			return;
 		}
 		entry.controller.abort("Cancelled by user");
-		this.activeStreams.delete(streamId);
+		this.untrackStream(streamId);
 		this.approvalGate.abortStream(streamId);
 		res.json({ cancelled: true });
 	}
@@ -878,5 +975,5 @@ function composePromptForAgent(registered, pluginLevel, ctx) {
 const agents = toPlugin(AgentsPlugin);
 
 //#endregion
-export { AgentsPlugin };
+export { AgentsPlugin, agents };
 //# sourceMappingURL=agents.js.map