@agentuity/workbench 1.0.5 → 1.0.7

package/src/components/internal/workbench-provider.tsx

@@ -1,6 +1,6 @@
 import type { WorkbenchConfig } from '@agentuity/core/workbench';
 import type React from 'react';
-import { createContext, useCallback, useContext, useEffect, useState } from 'react';
+import { createContext, useCallback, useContext, useEffect, useRef, useState } from 'react';
 import { useAgentSchemas } from '../../hooks/useAgentSchemas';
 import { useLogger } from '../../hooks/useLogger';
 import { useWorkbenchWebsocket } from '../../hooks/useWorkbenchWebsocket';
@@ -19,6 +19,13 @@ export function useWorkbench() {
 	return context;
 }
 
+/**
+ * Callback to get authentication headers for workbench requests.
+ * Called before each API request with the request body.
+ * Should return headers like X-Agentuity-Workbench-Signature and X-Agentuity-Workbench-Timestamp.
+ */
+export type GetAuthHeaders = (body: string) => Promise<Record<string, string>>;
+
 interface WorkbenchProviderProps {
 	config: Omit<WorkbenchConfig, 'route'> & {
 		baseUrl?: string | null;
@@ -36,6 +43,12 @@ interface WorkbenchProviderProps {
 			post?: React.ReactNode;
 		};
 	};
+	/**
+	 * Optional callback to get authentication headers for requests to deployed agents.
+	 * Called before each API request with the request body (empty string for GET/DELETE).
+	 * Useful for signature-based authentication in production deployments.
+	 */
+	getAuthHeaders?: GetAuthHeaders;
 }
 
 export function WorkbenchProvider({
@@ -47,9 +60,14 @@ export function WorkbenchProvider({
 	},
 	children,
 	portals,
+	getAuthHeaders,
 }: WorkbenchProviderProps) {
 	const logger = useLogger('WorkbenchProvider');
 
+	// Use ref for getAuthHeaders to prevent re-render loops when the callback changes identity
+	const getAuthHeadersRef = useRef(getAuthHeaders);
+	getAuthHeadersRef.current = getAuthHeaders;
+
 	// localStorage utilities scoped by project
 	const getStorageKey = useCallback(
 		(key: string) =>
@@ -77,26 +95,36 @@ export function WorkbenchProvider({
 		}
 	}, [getStorageKey]);
 
+	// Thread IDs are stored per baseUrl to avoid signature mismatch between environments
+	// (local signs with 'agentuity', cloud signs with AGENTUITY_SDK_KEY)
+	const getThreadStorageKey = useCallback(() => {
+		// Use a hash of the baseUrl to create unique storage per endpoint
+		const url = config.baseUrl ?? 'local';
+		const urlHash =
+			typeof url === 'string' ? btoa(encodeURIComponent(url)).slice(0, 16) : 'local';
+		return getStorageKey(`thread-id-${urlHash}`);
+	}, [getStorageKey, config.baseUrl]);
+
 	const saveThreadId = useCallback(
 		(threadId: string) => {
 			try {
-				localStorage.setItem(getStorageKey('thread-id'), threadId);
+				localStorage.setItem(getThreadStorageKey(), threadId);
 			} catch (error) {
 				logger.warn('Failed to save thread id to localStorage:', error);
 			}
 		},
-		[getStorageKey]
+		[getThreadStorageKey]
 	);
 
 	const loadThreadId = useCallback((): string | null => {
 		try {
-			return localStorage.getItem(getStorageKey('thread-id'));
+			return localStorage.getItem(getThreadStorageKey());
 		} catch (error) {
 			logger.warn('Failed to load thread id from localStorage:', error);
 
 			return null;
 		}
-	}, [getStorageKey]);
+	}, [getThreadStorageKey]);
 
 	const applyThreadIdHeader = useCallback(
 		(headers: Record<string, string>) => {
@@ -128,8 +156,51 @@ export function WorkbenchProvider({
 	// Config values
 	const baseUrl = config.baseUrl === undefined ? defaultBaseUrl : config.baseUrl;
 	const apiKey = config.apiKey;
+	const configHeaders = config.headers;
 	const isBaseUrlNull = config.baseUrl === null;
 
+	// Helper to build request headers with config headers, auth, and thread ID
+	const buildRequestHeaders = useCallback(
+		(additionalHeaders?: Record<string, string>): Record<string, string> => {
+			const headers: Record<string, string> = {
+				...(configHeaders || {}),
+				...(additionalHeaders || {}),
+			};
+
+			if (apiKey) {
+				headers.Authorization = `Bearer ${apiKey}`;
+			}
+
+			applyThreadIdHeader(headers);
+
+			return headers;
+		},
+		[configHeaders, apiKey, applyThreadIdHeader]
+	);
+
+	// Async helper to get request headers including auth headers from callback
+	const getRequestHeaders = useCallback(
+		async (
+			body: string,
+			additionalHeaders?: Record<string, string>
+		): Promise<Record<string, string>> => {
+			const headers = buildRequestHeaders(additionalHeaders);
+
+			// Call getAuthHeaders callback if provided (use ref to avoid re-render loops)
+			if (getAuthHeadersRef.current) {
+				try {
+					const authHeaders = await getAuthHeadersRef.current(body);
+					Object.assign(headers, authHeaders);
+				} catch (error) {
+					logger.warn('Failed to get auth headers:', error);
+				}
+			}
+
+			return headers;
+		},
+		[buildRequestHeaders, logger]
+	);
+
 	// Log baseUrl state
 	useEffect(() => {
 		if (isBaseUrlNull) {
@@ -139,14 +210,19 @@ export function WorkbenchProvider({
 		}
 	}, [isBaseUrlNull, baseUrl, logger]);
 
-	// Set disconnected status if baseUrl is null
+	// Set connection status based on baseUrl availability
+	// In cloud mode, we don't have websocket so we set connected when baseUrl is available
 	useEffect(() => {
 		if (isBaseUrlNull) {
 			logger.debug('🔌 Setting connection status to disconnected (baseUrl is null)');
-
 			setConnectionStatus('disconnected');
+		} else if (env.cloud) {
+			// In cloud mode, we're "connected" as soon as we have a baseUrl
+			// (no websocket to wait for)
+			logger.debug('🔌 Setting connection status to connected (cloud mode with baseUrl)');
+			setConnectionStatus('connected');
 		}
-	}, [isBaseUrlNull, logger]);
+	}, [isBaseUrlNull, env.cloud, logger]);
 
 	useEffect(() => {
 		if (isBaseUrlNull) {
@@ -160,24 +236,31 @@ export function WorkbenchProvider({
 		error: schemasError,
 		refetch: refetchSchemas,
 	} = useAgentSchemas({
-		baseUrl,
+		baseUrl: baseUrl ?? undefined,
 		apiKey,
+		headers: configHeaders,
 		enabled: !isBaseUrlNull,
+		getAuthHeaders,
 	});
 
 	// WebSocket connection for dev server restart detection
-	const wsBaseUrl = isBaseUrlNull ? undefined : baseUrl;
+	// Only enable for local dev - deployed agents don't have websocket endpoints
+	const wsEnabled = !isBaseUrlNull && !env.cloud;
+	const wsBaseUrl = wsEnabled && baseUrl ? baseUrl : undefined;
 
 	useEffect(() => {
 		if (isBaseUrlNull) {
 			logger.debug('🔌 WebSocket connection disabled (baseUrl is null)');
+		} else if (env.cloud) {
+			logger.debug('🔌 WebSocket connection disabled (cloud mode)');
 		}
-	}, [isBaseUrlNull, logger]);
+	}, [isBaseUrlNull, env.cloud, logger]);
 
 	const { connected } = useWorkbenchWebsocket({
-		enabled: !isBaseUrlNull,
+		enabled: wsEnabled,
 		baseUrl: wsBaseUrl,
 		apiKey,
+		headers: configHeaders,
 		onConnect: () => {
 			setConnectionStatus('connected');
 			refetchSchemas();
@@ -196,10 +279,12 @@ export function WorkbenchProvider({
 	});
 
 	useEffect(() => {
-		if (!isBaseUrlNull && !connected && connectionStatus !== 'restarting') {
+		// In cloud mode, websocket is disabled so we stay 'connected' (no live connection tracking)
+		// In local mode, track the websocket connection status
+		if (!isBaseUrlNull && !env.cloud && !connected && connectionStatus !== 'restarting') {
 			setConnectionStatus('disconnected');
 		}
-	}, [connected, connectionStatus, isBaseUrlNull]);
+	}, [connected, connectionStatus, isBaseUrlNull, env.cloud]);
 
 	// Convert schema data to Agent format, no fallback
 	const agents = schemaData?.agents;
@@ -228,18 +313,11 @@ export function WorkbenchProvider({
 			}
 
 			try {
-				const headers: Record<string, string> = {};
-
-				if (apiKey) {
-					headers.Authorization = `Bearer ${apiKey}`;
-				}
-
-				applyThreadIdHeader(headers);
-
 				const url = `${baseUrl}/_agentuity/workbench/state?agentId=${encodeURIComponent(agentId)}`;
 
 				logger.debug('📡 Fetching state for agent:', agentId);
 
+				const headers = await getRequestHeaders('');
 				const response = await fetch(url, {
 					method: 'GET',
 					headers,
@@ -295,7 +373,7 @@ export function WorkbenchProvider({
 				setMessages([]);
 			}
 		},
-		[baseUrl, apiKey, logger, applyThreadIdHeader, persistThreadIdFromResponse]
+		[baseUrl, logger, getRequestHeaders, persistThreadIdFromResponse]
 	);
 
 	// Set initial agent selection
@@ -369,6 +447,31 @@ export function WorkbenchProvider({
 		}
 	}, [agents, selectedAgent, loadSelectedAgent, saveSelectedAgent, logger, fetchAgentState]);
 
+	// Validate selected agent still exists when agents list changes (e.g., switching local ↔ cloud)
+	useEffect(() => {
+		if (!agents || Object.keys(agents).length === 0 || !selectedAgent) return;
+
+		const agentExists = Object.values(agents).some(
+			(agent) => agent.metadata.agentId === selectedAgent
+		);
+
+		if (!agentExists) {
+			logger.debug('⚠️ Selected agent no longer exists, falling back to first agent');
+
+			const sortedAgents = Object.values(agents).sort((a, b) =>
+				a.metadata.name.localeCompare(b.metadata.name)
+			);
+
+			const firstAgent = sortedAgents[0];
+
+			if (firstAgent) {
+				setSelectedAgent(firstAgent.metadata.agentId);
+				saveSelectedAgent(firstAgent.metadata.agentId);
+				fetchAgentState(firstAgent.metadata.agentId);
+			}
+		}
+	}, [agents, selectedAgent, logger, saveSelectedAgent, fetchAgentState]);
+
 	const submitMessage = async (value: string, _mode: 'text' | 'form' = 'text') => {
 		if (!selectedAgent) return;
 
@@ -456,16 +559,6 @@ export function WorkbenchProvider({
 
 			logger.debug('🌐 About to make API call...');
 
-			const headers: Record<string, string> = {
-				'Content-Type': 'application/json',
-			};
-
-			if (apiKey) {
-				headers.Authorization = `Bearer ${apiKey}`;
-			}
-
-			applyThreadIdHeader(headers);
-
 			const startTime = performance.now();
 
 			try {
@@ -473,13 +566,17 @@ export function WorkbenchProvider({
 					agentId: selectedAgent,
 					input: parsedInput,
 				};
+				const requestBody = JSON.stringify(requestPayload);
 
 				logger.debug('📤 API Request payload:', requestPayload);
 
+				const headers = await getRequestHeaders(requestBody, {
+					'Content-Type': 'application/json',
+				});
 				const response = await fetch(`${baseUrl}/_agentuity/workbench/execute`, {
 					method: 'POST',
 					headers,
-					body: JSON.stringify(requestPayload),
+					body: requestBody,
 					credentials: 'include',
 				});
 
@@ -641,19 +738,8 @@ export function WorkbenchProvider({
 
 		try {
 			const url = `${baseUrl}/_agentuity/workbench/sample?agentId=${encodeURIComponent(agentId)}`;
-			const headers: HeadersInit = {
-				'Content-Type': 'application/json',
-			};
-
-			if (apiKey) {
-				headers.Authorization = `Bearer ${apiKey}`;
-			}
-
-			// Keep thread id stable across workbench endpoints.
-			if (typeof headers === 'object' && headers && !Array.isArray(headers)) {
-				applyThreadIdHeader(headers as Record<string, string>);
-			}
 
+			const headers = await getRequestHeaders('', { 'Content-Type': 'application/json' });
 			const response = await fetch(url, {
 				method: 'GET',
 				headers,
@@ -713,15 +799,8 @@ export function WorkbenchProvider({
 			}
 
 			try {
-				const headers: Record<string, string> = {};
-
-				if (apiKey) {
-					headers.Authorization = `Bearer ${apiKey}`;
-				}
-
-				applyThreadIdHeader(headers);
-
 				const url = `${baseUrl}/_agentuity/workbench/state?agentId=${encodeURIComponent(agentId)}`;
+				const headers = await getRequestHeaders('');
 				const response = await fetch(url, {
 					method: 'DELETE',
 					headers,
@@ -741,7 +820,7 @@ export function WorkbenchProvider({
 				logger.debug('⚠️ Error clearing state:', error);
 			}
 		},
-		[baseUrl, apiKey, logger, applyThreadIdHeader, persistThreadIdFromResponse]
+		[baseUrl, logger, getRequestHeaders, persistThreadIdFromResponse]
 	);
 
 	const contextValue: WorkbenchContextType = {

package/src/hooks/useAgentSchemas.ts

@@ -1,5 +1,6 @@
 import type { JSONSchema7 } from 'ai';
-import { useCallback, useEffect, useState } from 'react';
+import { useCallback, useEffect, useRef, useState } from 'react';
+import type { GetAuthHeaders } from '../components/internal/workbench-provider';
 import { useLogger } from './useLogger';
 
 export interface AgentSchema {
@@ -36,6 +37,8 @@ export interface UseAgentSchemasOptions {
 	apiKey?: string;
 	baseUrl?: string;
 	enabled?: boolean;
+	headers?: Record<string, string>;
+	getAuthHeaders?: GetAuthHeaders;
 }
 
 export interface UseAgentSchemasResult {
@@ -66,13 +69,17 @@ export interface UseAgentSchemasResult {
  * ```
  */
 export function useAgentSchemas(options: UseAgentSchemasOptions = {}): UseAgentSchemasResult {
-	const { baseUrl = '', apiKey, enabled = true } = options;
+	const { baseUrl = '', apiKey, enabled = true, headers: configHeaders, getAuthHeaders } = options;
 
 	const logger = useLogger('useAgentSchemas');
 	const [data, setData] = useState<AgentSchemasResponse | null>(null);
 	const [isLoading, setIsLoading] = useState(false);
 	const [error, setError] = useState<Error | null>(null);
 
+	// Use ref to avoid re-render loops when getAuthHeaders changes
+	const getAuthHeadersRef = useRef(getAuthHeaders);
+	getAuthHeadersRef.current = getAuthHeaders;
+
 	const fetchSchemas = useCallback(async () => {
 		if (!enabled) return;
 
@@ -81,7 +88,8 @@ export function useAgentSchemas(options: UseAgentSchemasOptions = {}): UseAgentS
 
 		try {
 			const url = `${baseUrl}/_agentuity/workbench/metadata.json`;
-			const headers: HeadersInit = {
+			const headers: Record<string, string> = {
+				...(configHeaders || {}),
 				'Content-Type': 'application/json',
 			};
 
@@ -89,6 +97,16 @@ export function useAgentSchemas(options: UseAgentSchemasOptions = {}): UseAgentS
 				headers.Authorization = `Bearer ${apiKey}`;
 			}
 
+			// Get auth headers if callback is provided
+			if (getAuthHeadersRef.current) {
+				try {
+					const authHeaders = await getAuthHeadersRef.current('');
+					Object.assign(headers, authHeaders);
+				} catch (err) {
+					logger.warn('Failed to get auth headers:', err);
+				}
+			}
+
 			const response = await fetch(url, {
 				method: 'GET',
 				headers,
@@ -131,7 +149,7 @@ export function useAgentSchemas(options: UseAgentSchemasOptions = {}): UseAgentS
 		} finally {
 			setIsLoading(false);
 		}
-	}, [baseUrl, apiKey, enabled]);
+	}, [baseUrl, apiKey, enabled, configHeaders, logger]);
 
 	const refetch = useCallback(() => {
 		void fetchSchemas();

package/src/hooks/useLogger.ts

@@ -1,4 +1,4 @@
-import { useCallback } from 'react';
+import { useMemo } from 'react';
 
 type LogLevel = 'debug' | 'info' | 'warn' | 'error';
 
@@ -39,8 +39,11 @@ const shouldLog = (messageLevel: LogLevel): boolean => {
 };
 
 export function useLogger(component?: string): Logger {
-	const createLogFunction = useCallback(
-		(level: LogLevel) =>
+	// Memoize the entire logger object to prevent re-render loops
+	// when logger is used as a dependency in useCallback/useEffect
+	return useMemo(() => {
+		const createLogFunction =
+			(level: LogLevel) =>
 			(...args: unknown[]) => {
 				if (!shouldLog(level)) {
 					return;
@@ -50,14 +53,13 @@ export function useLogger(component?: string): Logger {
 				const consoleFn = console[level] || console.log;
 
 				consoleFn(prefix, ...args);
-			},
-		[component]
-	);
-
-	return {
-		debug: createLogFunction('debug'),
-		info: createLogFunction('info'),
-		warn: createLogFunction('warn'),
-		error: createLogFunction('error'),
-	};
+			};
+
+		return {
+			debug: createLogFunction('debug'),
+			info: createLogFunction('info'),
+			warn: createLogFunction('warn'),
+			error: createLogFunction('error'),
+		};
+	}, [component]);
 }

package/src/hooks/useWorkbenchWebsocket.ts

@@ -88,6 +88,7 @@ export interface UseWorkbenchWebsocketOptions {
 	apiKey?: string;
 	baseUrl?: string;
 	enabled?: boolean;
+	headers?: Record<string, string>;
 	onAlive?: () => void;
 	onConnect?: () => void;
 	onReconnect?: () => void;
@@ -102,7 +103,7 @@ export interface UseWorkbenchWebsocketResult {
 export function useWorkbenchWebsocket(
 	options: UseWorkbenchWebsocketOptions = {}
 ): UseWorkbenchWebsocketResult {
-	const { baseUrl, apiKey, onConnect, onReconnect, onAlive, onRestarting } = options;
+	const { baseUrl, apiKey, headers, onConnect, onReconnect, onAlive, onRestarting } = options;
 
 	const [connected, setConnected] = useState(false);
 	const [error, setError] = useState<Error | null>(null);
@@ -138,8 +139,20 @@ export function useWorkbenchWebsocket(
 			url.searchParams.set('apiKey', apiKey);
 		}
 
+		// Pass any manual headers as query params (WebSocket can't use custom headers in browser)
+		if (headers) {
+			const signature = headers['X-Agentuity-Workbench-Signature'];
+			const timestamp = headers['X-Agentuity-Workbench-Timestamp'];
+			if (signature) {
+				url.searchParams.set('signature', signature);
+			}
+			if (timestamp) {
+				url.searchParams.set('timestamp', timestamp);
+			}
+		}
+
 		return url.toString();
-	}, [baseUrl, apiKey]);
+	}, [baseUrl, apiKey, headers]);
 
 	const connect = useCallback(() => {
 		if (manualClose.current || !options.enabled) {

package/src/index.ts

@@ -2,7 +2,11 @@ export { default as App } from './components/App';
 export { Chat } from './components/internal/chat';
 export { StatusIndicator } from './components/internal/header';
 export { Schema } from './components/internal/schema';
-export { useWorkbench, WorkbenchProvider } from './components/internal/workbench-provider';
+export {
+	useWorkbench,
+	WorkbenchProvider,
+	type GetAuthHeaders,
+} from './components/internal/workbench-provider';
 export type { WorkbenchInstance } from './types';
 export type { ConnectionStatus, WorkbenchMessage } from './types/config';
 export { createWorkbench } from './workbench';