@btst/stack 2.2.0 → 2.4.0

package/src/plugins/ai-chat/api/plugin.ts

@@ -17,6 +17,10 @@ import {
 } from "../schemas";
 import type { Conversation, ConversationWithMessages, Message } from "../types";
 import { getAllConversations, getConversationById } from "./getters";
+import {
+	BUILT_IN_PAGE_TOOL_ROUTE_ALLOWLIST,
+	BUILT_IN_PAGE_TOOL_SCHEMAS,
+} from "./page-tools";
 
 /**
  * Context passed to AI Chat API hooks
@@ -98,6 +102,23 @@ export interface AiChatBackendHooks {
 		context: ChatApiContext,
 	) => Promise<boolean> | boolean;
 
+	/**
+	 * Called after the structural routeName/allowlist validation, with the list
+	 * of tool names that passed. Return a filtered subset to further restrict
+	 * which tools the LLM sees, or return [] to suppress all page tools.
+	 * Throw an Error to abort the entire chat request with a 403 response.
+	 * Not called when no tools passed the structural validation step.
+	 *
+	 * @param toolNames - Names that passed the routeName allowlist check
+	 * @param routeName - routeName claimed by the request (may be undefined)
+	 * @param context   - Full request context (headers, body, etc.)
+	 */
+	onBeforeToolsActivated?: (
+		toolNames: string[],
+		routeName: string | undefined,
+		context: ChatApiContext,
+	) => Promise<string[]> | string[];
+
 	// ============== Lifecycle Hooks ==============
 
 	/**
@@ -232,6 +253,32 @@ export type AiChatMode = "authenticated" | "public";
 /**
  * Configuration for AI Chat backend plugin
  */
+/**
+ * Extracts only the literal (non-index-signature) keys from a type.
+ * For `Record<string, T>` this resolves to `never`, so collision checks are
+ * skipped when the tools map is typed with a broad string index.
+ */
+type KnownKeys<T> = {
+	[K in keyof T]: string extends K ? never : K;
+}[keyof T];
+
+/**
+ * Ensures `TClientTools` has no keys that are also literal keys in `TTools`.
+ * Colliding keys are mapped to `never`, which produces a compile-time error
+ * at the point of the duplicate key. When `TTools` uses a string index
+ * signature the check is skipped to avoid false positives.
+ */
+type NoKeyCollision<
+	TTools,
+	TClientTools extends Record<string, Tool>,
+> = KnownKeys<TTools> & keyof TClientTools extends never
+	? TClientTools
+	: {
+			[K in keyof TClientTools]: K extends KnownKeys<TTools>
+				? never // duplicate of a server-side tool — remove from clientToolSchemas
+				: TClientTools[K];
+		};
+
 export interface AiChatBackendConfig {
 	/**
 	 * The language model to use for chat completions.
@@ -269,6 +316,31 @@ export interface AiChatBackendConfig {
 	 */
 	tools?: Record<string, Tool>;
 
+	/**
+	 * Enable route-aware page tools.
+	 * When true, the server will include tool schemas for client-side page tools
+	 * (e.g. fillBlogForm, updatePageLayers) based on the availableTools list
+	 * sent with each request.
+	 * @default false
+	 */
+	enablePageTools?: boolean;
+
+	/**
+	 * Custom client-side tool schemas for non-BTST pages.
+	 * Merged with built-in page tool schemas (fillBlogForm, updatePageLayers).
+	 * Only included when enablePageTools is true and the tool name appears in
+	 * the availableTools list sent with the request.
+	 *
+	 * @example
+	 * clientToolSchemas: {
+	 *   addToCart: tool({
+	 *     description: "Add current product to cart",
+	 *     parameters: z.object({ quantity: z.number().int().min(1) }),
+	 *   }),
+	 * }
+	 */
+	clientToolSchemas?: Record<string, Tool>;
+
 	/**
 	 * Optional hooks for customizing plugin behavior
 	 */
@@ -282,7 +354,15 @@ export interface AiChatBackendConfig {
  *
  * @param config - Configuration including model, tools, and optional hooks
  */
-export const aiChatBackendPlugin = (config: AiChatBackendConfig) =>
+export const aiChatBackendPlugin = <
+	TTools extends Record<string, Tool> = Record<never, Tool>,
+	TClientTools extends Record<string, Tool> = Record<never, Tool>,
+>(
+	config: Omit<AiChatBackendConfig, "tools" | "clientToolSchemas"> & {
+		tools?: TTools;
+		clientToolSchemas?: NoKeyCollision<TTools, TClientTools>;
+	},
+) =>
 	defineBackendPlugin({
 		name: "ai-chat",
 		// Always include db schema - in public mode we just don't use it
@@ -350,7 +430,13 @@ export const aiChatBackendPlugin = (config: AiChatBackendConfig) =>
 					body: chatRequestSchema,
 				},
 				async (ctx) => {
-					const { messages: rawMessages, conversationId } = ctx.body;
+					const {
+						messages: rawMessages,
+						conversationId,
+						pageContext,
+						availableTools,
+						routeName,
+					} = ctx.body;
 					const uiMessages = rawMessages as UIMessage[];
 
 					const context: ChatApiContext = {
@@ -388,22 +474,107 @@ export const aiChatBackendPlugin = (config: AiChatBackendConfig) =>
 						// Convert UIMessages to CoreMessages for streamText
 						const modelMessages = convertToModelMessages(uiMessages);
 
-						// Add system prompt if configured
-						const messagesWithSystem = config.systemPrompt
+						// Build system prompt: base config + optional page context
+						const pageContextContent =
+							pageContext && pageContext.trim()
+								? `\n\nCurrent page context:\n${pageContext}`
+								: "";
+						const systemContent = config.systemPrompt
+							? `${config.systemPrompt}${pageContextContent}`
+							: pageContextContent || undefined;
+
+						const messagesWithSystem = systemContent
 							? [
-									{ role: "system" as const, content: config.systemPrompt },
+									{ role: "system" as const, content: systemContent },
 									...modelMessages,
 								]
 							: modelMessages;
 
+						// Merge page tool schemas when enablePageTools is on.
+						// Built-in schemas are only included when the request's routeName is in
+						// the tool's allowlist — this prevents a page from claiming tools that
+						// are intended for a different route (e.g. requesting updatePageLayers
+						// from a blog page). Consumer clientToolSchemas are trusted as-is.
+						const activePageTools: Record<string, Tool> =
+							config.enablePageTools &&
+							availableTools &&
+							availableTools.length > 0
+								? (() => {
+										const consumerSchemas: Record<string, Tool> =
+											(config.clientToolSchemas as Record<string, Tool>) ?? {};
+										return Object.fromEntries(
+											availableTools
+												.filter((name) => {
+													// Built-in tool: require routeName to be in its allowlist
+													if (name in BUILT_IN_PAGE_TOOL_SCHEMAS) {
+														const allowed =
+															BUILT_IN_PAGE_TOOL_ROUTE_ALLOWLIST[name];
+														return (
+															allowed &&
+															routeName &&
+															allowed.includes(routeName)
+														);
+													}
+													// Consumer-defined tool: allow if schema is registered
+													return name in consumerSchemas;
+												})
+												.map((name) => {
+													const schema =
+														BUILT_IN_PAGE_TOOL_SCHEMAS[name] ??
+														consumerSchemas[name]!;
+													return [name, schema];
+												}),
+										);
+									})()
+								: {};
+
+						// Consumer hook: user-level tool authorization.
+						// Runs after the structural routeName allowlist check.
+						// A thrown Error is caught and returned as a 403 response,
+						// consistent with how onBeforeChat handles return false → 403.
+						if (
+							config.hooks?.onBeforeToolsActivated &&
+							Object.keys(activePageTools).length > 0
+						) {
+							try {
+								const allowed = await config.hooks.onBeforeToolsActivated(
+									Object.keys(activePageTools),
+									routeName,
+									context,
+								);
+								const allowedSet = new Set(allowed);
+								for (const key of Object.keys(activePageTools)) {
+									if (!allowedSet.has(key)) {
+										delete activePageTools[key];
+									}
+								}
+							} catch (hookError) {
+								throw ctx.error(403, {
+									message:
+										hookError instanceof Error
+											? hookError.message
+											: "Unauthorized: Tool activation denied",
+								});
+							}
+						}
+
+						// Page tools are layered under server-side tools so that a
+						// clientToolSchemas entry with the same name as a tool in
+						// config.tools never silently drops its `execute` function.
+						// Server-side tools always win on collision.
+						const mergedTools =
+							Object.keys(activePageTools).length > 0
+								? { ...activePageTools, ...config.tools }
+								: config.tools;
+
 						// PUBLIC MODE: Stream without persistence
 						if (isPublicMode) {
 							const result = streamText({
 								model: config.model,
 								messages: messagesWithSystem,
-								tools: config.tools,
+								tools: mergedTools,
 								// Enable multi-step tool calls if tools are configured
-								...(config.tools ? { stopWhen: stepCountIs(5) } : {}),
+								...(mergedTools ? { stopWhen: stepCountIs(5) } : {}),
 							});
 
 							return result.toUIMessageStreamResponse({
@@ -557,9 +728,9 @@ export const aiChatBackendPlugin = (config: AiChatBackendConfig) =>
 						const result = streamText({
 							model: config.model,
 							messages: messagesWithSystem,
-							tools: config.tools,
+							tools: mergedTools,
 							// Enable multi-step tool calls if tools are configured
-							...(config.tools ? { stopWhen: stepCountIs(5) } : {}),
+							...(mergedTools ? { stopWhen: stepCountIs(5) } : {}),
 							onFinish: async (completion: { text: string }) => {
 								// Wrap in try-catch since this runs after the response is sent
 								// and errors would otherwise become unhandled promise rejections

package/src/plugins/ai-chat/client/components/chat-input.tsx

@@ -260,14 +260,14 @@ export function ChatInput({
 				)}
 
 				{/* Text Input */}
-				<div className="relative flex-1">
+				<div className="relative flex-1 min-w-0">
 					<Textarea
 						value={input}
 						onChange={handleInputChange}
 						onKeyDown={handleKeyDown}
 						placeholder={placeholder || localization.CHAT_PLACEHOLDER}
 						className={cn(
-							"resize-none pr-12",
+							"resize-none pr-12 max-w-full",
 							isCompact
 								? "min-h-[40px] max-h-[120px] py-2"
 								: "min-h-[50px] max-h-[200px] py-3",

package/src/plugins/ai-chat/client/components/chat-interface.tsx

@@ -7,7 +7,11 @@ import { ChatMessage } from "./chat-message";
 import { ChatInput, type AttachedFile } from "./chat-input";
 import { StackAttribution } from "@workspace/ui/components/stack-attribution";
 import { ScrollArea } from "@workspace/ui/components/scroll-area";
-import { DefaultChatTransport, type UIMessage } from "ai";
+import {
+	DefaultChatTransport,
+	lastAssistantMessageIsCompleteWithToolCalls,
+	type UIMessage,
+} from "ai";
 import { cn } from "@workspace/ui/lib/utils";
 import { usePluginOverrides, useBasePath } from "@btst/stack/context";
 import type { AiChatPluginOverrides } from "../overrides";
@@ -20,6 +24,7 @@ import {
 	useConversations,
 	type SerializedConversation,
 } from "../hooks/chat-hooks";
+import { usePageAIContext } from "../context/page-ai-context";
 
 interface ChatInterfaceProps {
 	apiPath?: string;
@@ -56,6 +61,9 @@ export function ChatInterface({
 	const basePath = useBasePath();
 	const isPublicMode = mode === "public";
 
+	// Read page AI context registered by the current page
+	const pageAIContext = usePageAIContext();
+
 	const localization = { ...AI_CHAT_LOCALIZATION, ...customLocalization };
 	const queryClient = useQueryClient();
 
@@ -126,6 +134,13 @@ export function ChatInterface({
 			!initialMessages || initialMessages.length === 0,
 	);
 
+	// Ref to always have the latest pageAIContext in the transport callback
+	// without recreating the transport on every context change
+	const pageAIContextRef = useRef(pageAIContext);
+	useEffect(() => {
+		pageAIContextRef.current = pageAIContext;
+	}, [pageAIContext]);
+
 	// Memoize the transport to prevent recreation on every render
 	const transport = useMemo(
 		() =>
@@ -135,8 +150,21 @@ export function ChatInterface({
 				body: isPublicMode
 					? undefined
 					: () => ({ conversationId: conversationIdRef.current }),
-				// Handle edit operations by using truncated messages from the ref
+				// Handle edit operations and inject page context
 				prepareSendMessagesRequest: ({ messages: hookMessages }) => {
+					const currentPageContext = pageAIContextRef.current;
+
+					// Build page context fields to include in every request
+					const pageContextBody = currentPageContext?.pageDescription
+						? {
+								pageContext: currentPageContext.pageDescription,
+								availableTools: Object.keys(
+									currentPageContext.clientTools ?? {},
+								),
+								routeName: currentPageContext.routeName,
+							}
+						: {};
+
 					// If we're in an edit operation, use the truncated messages + new user message
 					if (editMessagesRef.current !== null) {
 						const newUserMessage = hookMessages[hookMessages.length - 1];
@@ -150,6 +178,7 @@ export function ChatInterface({
 							body: {
 								messages: messagesToSend,
 								conversationId: conversationIdRef.current,
+								...pageContextBody,
 							},
 						};
 					}
@@ -158,6 +187,7 @@ export function ChatInterface({
 						body: {
 							messages: hookMessages,
 							conversationId: conversationIdRef.current,
+							...pageContextBody,
 						},
 					};
 				},
@@ -165,48 +195,99 @@ export function ChatInterface({
 		[apiPath, isPublicMode],
 	);
 
-	const { messages, sendMessage, status, error, setMessages, regenerate } =
-		useChat({
-			transport,
-			onError: (err) => {
-				console.error("useChat onError:", err);
-				// Reset first-message tracking if the send failed before a conversation was created.
-				// Without this, isFirstMessageSentRef stays true and the next successful send
-				// skips the "first message" navigation logic, corrupting the conversation flow.
-				if (!id && !hasNavigatedRef.current) {
-					isFirstMessageSentRef.current = false;
-				}
-			},
-			onFinish: async () => {
-				// In public mode, skip all persistence-related operations
-				if (isPublicMode) return;
+	// Use a ref so addToolOutput is always current inside the onToolCall closure
+	const addToolOutputRef = useRef<
+		ReturnType<typeof useChat>["addToolOutput"] | null
+	>(null);
 
-				// Invalidate conversation list to show new/updated conversations
-				await queryClient.invalidateQueries({
+	const {
+		messages,
+		sendMessage,
+		status,
+		error,
+		setMessages,
+		regenerate,
+		addToolOutput,
+	} = useChat({
+		transport,
+		// Automatically resubmit after all client-side tool results are provided
+		sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+		onToolCall: async ({ toolCall }) => {
+			// Dispatch client-side tool calls to the handler registered by the current page.
+			// In AI SDK v5, onToolCall returns void — addToolOutput must be called explicitly.
+			const toolName = toolCall.toolName;
+			const handler = pageAIContextRef.current?.clientTools?.[toolName];
+			if (handler) {
+				try {
+					const result = await handler(toolCall.input);
+					// No await — avoids potential deadlocks with sendAutomaticallyWhen
+					addToolOutputRef.current?.({
+						tool: toolName,
+						toolCallId: toolCall.toolCallId,
+						output: result,
+					});
+				} catch (err) {
+					addToolOutputRef.current?.({
+						tool: toolName,
+						toolCallId: toolCall.toolCallId,
+						state: "output-error",
+						errorText:
+							err instanceof Error ? err.message : "Tool execution failed",
+					});
+				}
+			} else {
+				// No handler found — this happens when the user navigates away while a
+				// tool-call response is streaming and the page context changes. Always
+				// call addToolOutput so sendAutomaticallyWhen can unblock; without this
+				// the conversation gets permanently stuck waiting for a missing output.
+				addToolOutputRef.current?.({
+					tool: toolName,
+					toolCallId: toolCall.toolCallId,
+					state: "output-error",
+					errorText: `No client-side handler registered for tool "${toolName}". The page context may have changed while the response was streaming.`,
+				});
+			}
+		},
+		onError: (err) => {
+			console.error("useChat onError:", err);
+			// Reset first-message tracking if the send failed before a conversation was created.
+			// Without this, isFirstMessageSentRef stays true and the next successful send
+			// skips the "first message" navigation logic, corrupting the conversation flow.
+			if (!id && !hasNavigatedRef.current) {
+				isFirstMessageSentRef.current = false;
+			}
+		},
+		onFinish: async () => {
+			// In public mode, skip all persistence-related operations
+			if (isPublicMode) return;
+
+			// Invalidate conversation list to show new/updated conversations
+			await queryClient.invalidateQueries({
+				queryKey: conversationsListQueryKey,
+			});
+
+			// If this was the first message on a new chat, update the URL without full navigation
+			// This avoids losing the in-memory messages during component remount
+			if (isFirstMessageSentRef.current && !id && !hasNavigatedRef.current) {
+				hasNavigatedRef.current = true;
+				// Wait for the invalidation to complete and refetch conversations
+				await queryClient.refetchQueries({
 					queryKey: conversationsListQueryKey,
 				});
-
-				// If this was the first message on a new chat, update the URL without full navigation
-				// This avoids losing the in-memory messages during component remount
-				if (isFirstMessageSentRef.current && !id && !hasNavigatedRef.current) {
-					hasNavigatedRef.current = true;
-					// Wait for the invalidation to complete and refetch conversations
-					await queryClient.refetchQueries({
-						queryKey: conversationsListQueryKey,
-					});
-					// Get the updated conversations from cache
-					const cachedConversations = queryClient.getQueryData<
-						SerializedConversation[]
-					>(conversationsListQueryKey);
-					if (cachedConversations && cachedConversations.length > 0) {
-						// The most recently updated conversation should be the one we just created
-						const newConversation = cachedConversations[0];
-						if (newConversation) {
-							// Update our local state
-							setCurrentConversationId(newConversation.id);
-							conversationIdRef.current = newConversation.id;
-							// Update URL without navigation to preserve in-memory messages
-							// Use replaceState to avoid adding to history stack
+				// Get the updated conversations from cache
+				const cachedConversations = queryClient.getQueryData<
+					SerializedConversation[]
+				>(conversationsListQueryKey);
+				if (cachedConversations && cachedConversations.length > 0) {
+					// The most recently updated conversation should be the one we just created
+					const newConversation = cachedConversations[0];
+					if (newConversation) {
+						// Update our local state
+						setCurrentConversationId(newConversation.id);
+						conversationIdRef.current = newConversation.id;
+						// Only update the URL in full-page mode; in widget mode the chat is
+						// embedded in another page and clobbering the URL is disruptive.
+						if (variant === "full") {
 							const newUrl = `${basePath}/chat/${newConversation.id}`;
 							if (typeof window !== "undefined") {
 								window.history.replaceState(
@@ -218,8 +299,14 @@ export function ChatInterface({
 						}
 					}
 				}
-			},
-		});
+			}
+		},
+	});
+
+	// Keep addToolOutputRef in sync so onToolCall always has the latest reference
+	useEffect(() => {
+		addToolOutputRef.current = addToolOutput;
+	}, [addToolOutput]);
 
 	// Load existing conversation messages when navigating to a conversation
 	useEffect(() => {
@@ -484,23 +571,32 @@ export function ChatInterface({
 						>
 							{messages.length === 0 ? (
 								<div className="flex flex-col h-full min-h-[300px]">
-									<div className="flex-1 flex items-center justify-center text-muted-foreground">
+									<div className="flex-1 flex items-center justify-center text-muted-foreground mb-4">
 										<p>{localization.CHAT_EMPTY_STATE}</p>
 									</div>
-									{chatSuggestions && chatSuggestions.length > 0 && (
-										<div className="flex flex-wrap justify-center gap-2 pb-4 max-w-md mx-auto">
-											{chatSuggestions.map((suggestion, index) => (
-												<button
-													key={index}
-													type="button"
-													onClick={() => setInput(suggestion)}
-													className="px-3 py-2 text-sm rounded-lg border border-border bg-background hover:bg-accent hover:text-accent-foreground transition-colors text-foreground"
-												>
-													{suggestion}
-												</button>
-											))}
-										</div>
-									)}
+									{(() => {
+										// Merge static suggestions from overrides with dynamic ones from page context.
+										// Page context suggestions appear first (most relevant to current page).
+										const pageSuggestions = pageAIContext?.suggestions ?? [];
+										const allSuggestions = [
+											...pageSuggestions,
+											...(chatSuggestions ?? []),
+										];
+										return allSuggestions.length > 0 ? (
+											<div className="flex flex-wrap justify-center gap-2 pb-4 max-w-md mx-auto">
+												{allSuggestions.map((suggestion, index) => (
+													<button
+														key={index}
+														type="button"
+														onClick={() => setInput(suggestion)}
+														className="px-3 py-2 text-sm rounded-lg border border-border bg-background hover:bg-accent hover:text-accent-foreground transition-colors text-foreground"
+													>
+														{suggestion}
+													</button>
+												))}
+											</div>
+										) : null;
+									})()}
 								</div>
 							) : (
 								messages.map((m, index) => (