@datalayer/agent-runtimes 0.0.8 → 0.0.10

package/lib/components/chat/components/base/ChatBase.js

@@ -26,6 +26,7 @@ import { Streamdown } from 'streamdown';
 import { PoweredByTag } from '../elements/PoweredByTag';
 import { requestAPI } from '../../handler';
 import { useChatStore } from '../../store/chatStore';
+import { useConversationStore } from '../../store/conversationStore';
 import { generateMessageId, createUserMessage, createAssistantMessage, } from '../../types/message';
 import { AGUIAdapter, A2AAdapter, VercelAIAdapter, ACPAdapter, } from '../../protocols';
 import { ToolCallDisplay } from '../display/ToolCallDisplay';
@@ -155,6 +156,7 @@ function useConfigQuery(enabled, configEndpoint, authToken) {
         },
         queryKey: ['models', configEndpoint || 'jupyter'],
         enabled,
+        retry: 1,
     });
 }
 /**
@@ -195,37 +197,114 @@ function useSkillsQuery(enabled, baseEndpoint, authToken) {
         queryKey: ['skills', baseEndpoint || 'jupyter'],
         enabled,
         staleTime: 5 * 60 * 1000, // 5 minutes
+        retry: 1,
     });
 }
+/**
+ * Format token count for compact display
+ */
+function formatTokenCount(tokens) {
+    if (tokens >= 1_000_000)
+        return `${(tokens / 1_000_000).toFixed(1)}M`;
+    if (tokens >= 1_000)
+        return `${(tokens / 1_000).toFixed(1)}K`;
+    return tokens.toString();
+}
+/**
+ * Derive the API base URL from a configEndpoint.
+ * configEndpoint may look like:
+ *   "http://host:port/api/v1/config"     (from agentRuntimeConfig)
+ *   "http://host:port/api/v1/configure"   (from Chat component)
+ * We strip the trailing path segment to get "http://host:port/api/v1".
+ */
+function getApiBaseFromConfig(configEndpoint) {
+    return configEndpoint.replace(/\/(config|configure)\/?$/, '');
+}
+/**
+ * Hook to poll agent context-snapshot from the backend.
+ * Returns cumulative token usage (input/output breakdown) tracked by the agent server.
+ * Uses the same endpoint as codeai: GET /api/v1/configure/agents/{agentId}/context-snapshot
+ */
+function useContextSnapshotQuery(enabled, configEndpoint, agentId, authToken) {
+    const queryClient = useContext(QueryClientContext);
+    if (!queryClient) {
+        return { data: undefined, isLoading: false, isError: false, error: null };
+    }
+    const snapshotUrl = configEndpoint && agentId
+        ? `${getApiBaseFromConfig(configEndpoint)}/configure/agents/${encodeURIComponent(agentId)}/context-snapshot`
+        : undefined;
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    const result = useQuery({
+        queryKey: ['context-snapshot-header', agentId, snapshotUrl],
+        queryFn: async () => {
+            if (!snapshotUrl) {
+                throw new Error('No context-snapshot URL available');
+            }
+            const headers = { 'Content-Type': 'application/json' };
+            if (authToken) {
+                headers['Authorization'] = `Bearer ${authToken}`;
+            }
+            const response = await fetch(snapshotUrl, { headers });
+            if (!response.ok) {
+                throw new Error(`Context snapshot fetch failed: ${response.statusText}`);
+            }
+            return response.json();
+        },
+        enabled: enabled && !!snapshotUrl,
+        // Poll every 10 seconds, but stop polling once the query has errored (e.g. runtime terminated)
+        refetchInterval: query => (query.state.status === 'error' ? false : 10_000),
+        refetchOnMount: 'always',
+        staleTime: 0,
+        retry: 1,
+    });
+    return result;
+}
 /**
  * ChatBase component - Universal chat panel supporting store, protocol, and custom modes
  */
-export function ChatBase({ title, showHeader = false, showLoadingIndicator = true, showErrors = true, showInput = true, showModelSelector = false, showToolsMenu = false, showSkillsMenu = false, codemodeEnabled = false, initialModel, initialMcpServers, initialSkills, className, loadingState, headerActions, 
+export function ChatBase({ title, showHeader = false, showTokenUsage = true, showLoadingIndicator = true, showErrors = true, showInput = true, showModelSelector = false, showToolsMenu = false, showSkillsMenu = false, codemodeEnabled = false, initialModel, availableModels, mcpServers, initialSkills, className, loadingState, headerActions, 
 // Mode selection
-useStore: useStoreMode = true, protocol, onSendMessage, enableStreaming = false, 
+useStore: useStoreMode = true, protocol: protocolProp, agentRuntimeConfig, onSendMessage, enableStreaming = false, 
 // Extended props
 brandIcon, avatarConfig, headerButtons, showPoweredBy = false, poweredByProps, emptyState, renderToolResult, footerContent, headerContent, children, borderRadius, backgroundColor, border, boxShadow, compact = false, placeholder, description = 'Start a conversation with the AI agent.', onStateUpdate, onNewChat, onClear, onMessagesChange, autoFocus = false, suggestions, submitOnSuggestionClick = true, hideMessagesAfterToolUI = false, focusTrigger, frontendTools, 
 // Identity/Authorization props
-onAuthorizationRequired, connectedIdentities, }) {
+onAuthorizationRequired, connectedIdentities, 
+// Conversation persistence
+runtimeId, historyEndpoint, historyAuthToken, }) {
+    // Convert agentRuntimeConfig to protocol if provided
+    const protocol = agentRuntimeConfig
+        ? {
+            type: agentRuntimeConfig.protocol || 'ag-ui',
+            endpoint: agentRuntimeConfig.url,
+            authToken: agentRuntimeConfig.authToken,
+            agentId: agentRuntimeConfig.agentId,
+            enableConfigQuery: true,
+            configEndpoint: `${agentRuntimeConfig.url}/api/v1/config`,
+        }
+        : protocolProp;
+    // If agentRuntimeConfig is provided, force protocol mode
+    const effectiveUseStoreMode = agentRuntimeConfig ? false : useStoreMode;
     // Check if QueryClientProvider is already available
     const existingQueryClient = useContext(QueryClientContext);
     // If no QueryClient is available, wrap with our internal provider
     if (!existingQueryClient) {
-        return (_jsx(QueryClientProvider, { client: internalQueryClient, children: _jsx(ChatBaseInner, { title: title, showHeader: showHeader, showLoadingIndicator: showLoadingIndicator, showErrors: showErrors, showInput: showInput, showModelSelector: showModelSelector, showToolsMenu: showToolsMenu, showSkillsMenu: showSkillsMenu, codemodeEnabled: codemodeEnabled, initialModel: initialModel, initialMcpServers: initialMcpServers, initialSkills: initialSkills, className: className, loadingState: loadingState, headerActions: headerActions, useStore: useStoreMode, protocol: protocol, onSendMessage: onSendMessage, enableStreaming: enableStreaming, brandIcon: brandIcon, avatarConfig: avatarConfig, headerButtons: headerButtons, showPoweredBy: showPoweredBy, poweredByProps: poweredByProps, emptyState: emptyState, renderToolResult: renderToolResult, footerContent: footerContent, headerContent: headerContent, children: children, borderRadius: borderRadius, backgroundColor: backgroundColor, border: border, boxShadow: boxShadow, compact: compact, placeholder: placeholder, description: description, onStateUpdate: onStateUpdate, onNewChat: onNewChat, onClear: onClear, onMessagesChange: onMessagesChange, autoFocus: autoFocus, suggestions: suggestions, submitOnSuggestionClick: submitOnSuggestionClick, hideMessagesAfterToolUI: hideMessagesAfterToolUI, focusTrigger: focusTrigger, frontendTools: frontendTools, onAuthorizationRequired: onAuthorizationRequired, connectedIdentities: connectedIdentities }) }));
+        return (_jsx(QueryClientProvider, { client: internalQueryClient, children: _jsx(ChatBaseInner, { title: title, showHeader: showHeader, showTokenUsage: showTokenUsage, showLoadingIndicator: showLoadingIndicator, showErrors: showErrors, showInput: showInput, showModelSelector: showModelSelector, showToolsMenu: showToolsMenu, showSkillsMenu: showSkillsMenu, codemodeEnabled: codemodeEnabled, initialModel: initialModel, availableModels: availableModels, mcpServers: mcpServers, initialSkills: initialSkills, className: className, loadingState: loadingState, headerActions: headerActions, useStore: effectiveUseStoreMode, protocol: protocol, onSendMessage: onSendMessage, enableStreaming: enableStreaming, brandIcon: brandIcon, avatarConfig: avatarConfig, headerButtons: headerButtons, showPoweredBy: showPoweredBy, poweredByProps: poweredByProps, emptyState: emptyState, renderToolResult: renderToolResult, footerContent: footerContent, headerContent: headerContent, children: children, borderRadius: borderRadius, backgroundColor: backgroundColor, border: border, boxShadow: boxShadow, compact: compact, placeholder: placeholder, description: description, onStateUpdate: onStateUpdate, onNewChat: onNewChat, onClear: onClear, onMessagesChange: onMessagesChange, autoFocus: autoFocus, suggestions: suggestions, submitOnSuggestionClick: submitOnSuggestionClick, hideMessagesAfterToolUI: hideMessagesAfterToolUI, focusTrigger: focusTrigger, frontendTools: frontendTools, onAuthorizationRequired: onAuthorizationRequired, connectedIdentities: connectedIdentities, runtimeId: runtimeId, historyEndpoint: historyEndpoint, historyAuthToken: historyAuthToken }) }));
     }
     // QueryClient already available, render inner component directly
-    return (_jsx(ChatBaseInner, { title: title, showHeader: showHeader, showLoadingIndicator: showLoadingIndicator, showErrors: showErrors, showInput: showInput, showModelSelector: showModelSelector, showToolsMenu: showToolsMenu, showSkillsMenu: showSkillsMenu, codemodeEnabled: codemodeEnabled, initialModel: initialModel, initialMcpServers: initialMcpServers, initialSkills: initialSkills, className: className, loadingState: loadingState, headerActions: headerActions, useStore: useStoreMode, protocol: protocol, onSendMessage: onSendMessage, enableStreaming: enableStreaming, brandIcon: brandIcon, avatarConfig: avatarConfig, headerButtons: headerButtons, showPoweredBy: showPoweredBy, poweredByProps: poweredByProps, emptyState: emptyState, renderToolResult: renderToolResult, footerContent: footerContent, headerContent: headerContent, children: children, borderRadius: borderRadius, backgroundColor: backgroundColor, border: border, boxShadow: boxShadow, compact: compact, placeholder: placeholder, description: description, onStateUpdate: onStateUpdate, onNewChat: onNewChat, onClear: onClear, onMessagesChange: onMessagesChange, autoFocus: autoFocus, suggestions: suggestions, submitOnSuggestionClick: submitOnSuggestionClick, hideMessagesAfterToolUI: hideMessagesAfterToolUI, focusTrigger: focusTrigger, frontendTools: frontendTools, onAuthorizationRequired: onAuthorizationRequired, connectedIdentities: connectedIdentities }));
+    return (_jsx(ChatBaseInner, { title: title, showHeader: showHeader, showTokenUsage: showTokenUsage, showLoadingIndicator: showLoadingIndicator, showErrors: showErrors, showInput: showInput, showModelSelector: showModelSelector, showToolsMenu: showToolsMenu, showSkillsMenu: showSkillsMenu, codemodeEnabled: codemodeEnabled, initialModel: initialModel, availableModels: availableModels, mcpServers: mcpServers, initialSkills: initialSkills, className: className, loadingState: loadingState, headerActions: headerActions, useStore: effectiveUseStoreMode, protocol: protocol, onSendMessage: onSendMessage, enableStreaming: enableStreaming, brandIcon: brandIcon, avatarConfig: avatarConfig, headerButtons: headerButtons, showPoweredBy: showPoweredBy, poweredByProps: poweredByProps, emptyState: emptyState, renderToolResult: renderToolResult, footerContent: footerContent, headerContent: headerContent, children: children, borderRadius: borderRadius, backgroundColor: backgroundColor, border: border, boxShadow: boxShadow, compact: compact, placeholder: placeholder, description: description, onStateUpdate: onStateUpdate, onNewChat: onNewChat, onClear: onClear, onMessagesChange: onMessagesChange, autoFocus: autoFocus, suggestions: suggestions, submitOnSuggestionClick: submitOnSuggestionClick, hideMessagesAfterToolUI: hideMessagesAfterToolUI, focusTrigger: focusTrigger, frontendTools: frontendTools, onAuthorizationRequired: onAuthorizationRequired, connectedIdentities: connectedIdentities, runtimeId: runtimeId, historyEndpoint: historyEndpoint, historyAuthToken: historyAuthToken }));
 }
 /**
  * Inner ChatBase component - contains all the actual logic
  */
-function ChatBaseInner({ title, showHeader = false, showLoadingIndicator = true, showErrors = true, showInput = true, showModelSelector = false, showToolsMenu = false, showSkillsMenu = false, codemodeEnabled = false, initialModel, initialMcpServers, initialSkills, className, loadingState, headerActions, 
+function ChatBaseInner({ title, showHeader = false, showTokenUsage = true, showLoadingIndicator = true, showErrors = true, showInput = true, showModelSelector = false, showToolsMenu = false, showSkillsMenu = false, codemodeEnabled = false, initialModel, availableModels, mcpServers, initialSkills, className, loadingState, headerActions, 
 // Mode selection
 useStore: useStoreMode = true, protocol, onSendMessage, enableStreaming = false, 
 // Extended props
 brandIcon, avatarConfig, headerButtons, showPoweredBy = false, poweredByProps, emptyState, renderToolResult, footerContent, headerContent, children, borderRadius, backgroundColor, border, boxShadow, compact = false, placeholder, description = 'Start a conversation with the AI agent.', onStateUpdate, onNewChat, onClear, onMessagesChange, autoFocus = false, suggestions, submitOnSuggestionClick = true, hideMessagesAfterToolUI = false, focusTrigger, frontendTools, 
 // Identity/Authorization props
-onAuthorizationRequired, connectedIdentities, }) {
+onAuthorizationRequired, connectedIdentities, 
+// Conversation persistence
+runtimeId, historyEndpoint, historyAuthToken, }) {
     // Ensure Primer's default portal has high z-index for ActionMenu overlays
     useHighZIndexPortal();
     // Store (optional for message persistence)
@@ -253,6 +332,10 @@ onAuthorizationRequired, connectedIdentities, }) {
     const configQuery = useConfigQuery(Boolean(protocol?.enableConfigQuery), protocol?.configEndpoint, protocol?.authToken);
     // Skills query (for protocols that support it)
     const skillsQuery = useSkillsQuery(Boolean(protocol?.enableConfigQuery) && showSkillsMenu, protocol?.configEndpoint, protocol?.authToken);
+    // Context snapshot query — polls agent-level cumulative token usage (input/output breakdown)
+    // Gated by showTokenUsage (not showHeader), so usage is visible even without a title bar.
+    const contextSnapshotQuery = useContextSnapshotQuery(Boolean(protocol?.enableConfigQuery) && showTokenUsage, protocol?.configEndpoint, protocol?.agentId, protocol?.authToken);
+    const agentUsage = contextSnapshotQuery.data;
     // Refs
     const adapterRef = useRef(null);
     const unsubscribeRef = useRef(null);
@@ -266,6 +349,12 @@ onAuthorizationRequired, connectedIdentities, }) {
     // (the array reference changes on every render even if contents are the same)
     const connectedIdentitiesRef = useRef(connectedIdentities);
     connectedIdentitiesRef.current = connectedIdentities;
+    const isServerSelected = useCallback((server) => {
+        if (!mcpServers)
+            return true;
+        const origin = server.isConfig === false ? 'catalog' : 'config';
+        return mcpServers.some(s => s.id === server.id && s.origin === origin);
+    }, [mcpServers]);
     // Auto-focus input on mount
     useEffect(() => {
         if (autoFocus && inputRef.current) {
@@ -326,18 +415,20 @@ onAuthorizationRequired, connectedIdentities, }) {
     }, [adjustTextareaHeight]);
     // Initialize model and tools when config is available
     useEffect(() => {
-        if (configQuery.data && !selectedModel) {
+        if ((configQuery.data || availableModels) && !selectedModel) {
+            // Use availableModels override if provided, otherwise use config models
+            const modelsList = availableModels || configQuery.data?.models || [];
             // Use initialModel if provided, otherwise select first available model
             if (initialModel) {
-                // Check if the initial model exists in the config
-                const modelExists = configQuery.data.models.some(m => m.id === initialModel);
+                // Check if the initial model exists in the models list
+                const modelExists = modelsList.some(m => m.id === initialModel);
                 if (modelExists) {
                     setSelectedModel(initialModel);
                 }
                 else {
                     // Fallback to first available model if initialModel not found
-                    const firstAvailableModel = configQuery.data.models.find(m => m.isAvailable !== false);
-                    const firstModel = firstAvailableModel || configQuery.data.models[0];
+                    const firstAvailableModel = modelsList.find(m => m.isAvailable !== false);
+                    const firstModel = firstAvailableModel || modelsList[0];
                     if (firstModel) {
                         setSelectedModel(firstModel.id);
                     }
@@ -345,24 +436,22 @@ onAuthorizationRequired, connectedIdentities, }) {
             }
             else {
                 // No initialModel provided, select first available model
-                const firstAvailableModel = configQuery.data.models.find(m => m.isAvailable !== false);
-                const firstModel = firstAvailableModel || configQuery.data.models[0];
+                const firstAvailableModel = modelsList.find(m => m.isAvailable !== false);
+                const firstModel = firstAvailableModel || modelsList[0];
                 if (firstModel) {
                     setSelectedModel(firstModel.id);
                 }
             }
-            const allToolIds = configQuery.data.builtinTools?.map(tool => tool.id) || [];
+            const allToolIds = configQuery.data?.builtinTools?.map(tool => tool.id) || [];
             setEnabledTools(allToolIds);
             // Initialize MCP server tools
-            if (configQuery.data.mcpServers) {
+            if (configQuery.data?.mcpServers) {
                 const newEnabledMcpTools = new Map();
                 for (const server of configQuery.data.mcpServers) {
                     if (server.isAvailable && server.enabled) {
-                        // If initialMcpServers is provided, only enable those servers
+                        // If mcpServers is provided, only enable those servers
                         // If not provided, enable all available servers
-                        const shouldEnableServer = initialMcpServers
-                            ? initialMcpServers.includes(server.id)
-                            : true;
+                        const shouldEnableServer = isServerSelected(server);
                         if (shouldEnableServer) {
                             const enabledToolNames = new Set(server.tools.filter(t => t.enabled).map(t => t.name));
                             newEnabledMcpTools.set(server.id, enabledToolNames);
@@ -372,7 +461,38 @@ onAuthorizationRequired, connectedIdentities, }) {
                 setEnabledMcpTools(newEnabledMcpTools);
             }
         }
-    }, [configQuery.data, selectedModel, initialModel, initialMcpServers]);
+    }, [
+        configQuery.data,
+        selectedModel,
+        initialModel,
+        availableModels,
+        mcpServers,
+        isServerSelected,
+    ]);
+    // Update enabled MCP servers when mcpServers prop changes (e.g., server removed)
+    useEffect(() => {
+        if (!configQuery.data?.mcpServers || !mcpServers)
+            return;
+        setEnabledMcpTools(prev => {
+            const newMap = new Map();
+            // Only keep servers that are in mcpServers
+            for (const server of configQuery.data?.mcpServers ?? []) {
+                if (isServerSelected(server) && prev.has(server.id)) {
+                    // Keep existing tool selection for this server
+                    newMap.set(server.id, prev.get(server.id));
+                }
+                else if (isServerSelected(server) &&
+                    server.isAvailable &&
+                    server.enabled) {
+                    // Newly added server - enable all tools
+                    const enabledToolNames = new Set(server.tools.filter(t => t.enabled).map(t => t.name));
+                    newMap.set(server.id, enabledToolNames);
+                }
+                // Servers not in mcpServers are excluded
+            }
+            return newMap;
+        });
+    }, [mcpServers, configQuery.data?.mcpServers, isServerSelected]);
     // Initialize enabled skills from initialSkills prop
     useEffect(() => {
         if (initialSkills && initialSkills.length > 0) {
@@ -425,13 +545,17 @@ onAuthorizationRequired, connectedIdentities, }) {
         setEnabledSkills(enable ? new Set(allSkillIds) : new Set());
     }, []);
     // Get all enabled MCP tool names (for sending with requests)
+    // Only counts tools from servers that are in mcpServers (if provided)
     const getEnabledMcpToolNames = useCallback(() => {
         const toolNames = [];
-        enabledMcpTools.forEach(tools => {
-            tools.forEach(toolName => toolNames.push(toolName));
+        enabledMcpTools.forEach((tools, serverId) => {
+            // Filter by mcpServers if provided
+            if (!mcpServers || mcpServers.some(s => s.id === serverId)) {
+                tools.forEach(toolName => toolNames.push(toolName));
+            }
         });
         return toolNames;
-    }, [enabledMcpTools]);
+    }, [enabledMcpTools, mcpServers]);
     // Get all enabled skill IDs (for sending with requests)
     const getEnabledSkillIds = useCallback(() => {
         return Array.from(enabledSkills);
@@ -445,6 +569,95 @@ onAuthorizationRequired, connectedIdentities, }) {
             }
         }
     }, [useStoreMode]);
+    // Track previous runtimeId to detect changes
+    const prevRuntimeIdRef = useRef(undefined);
+    // Clear displayItems and load messages when runtimeId changes
+    // This ensures each agent space has isolated conversation history
+    useEffect(() => {
+        // If runtimeId changed, clear displayItems first
+        if (runtimeId !== prevRuntimeIdRef.current) {
+            prevRuntimeIdRef.current = runtimeId;
+            // Clear current display items when switching runtime
+            setDisplayItems([]);
+            toolCallsRef.current.clear();
+            // If no runtimeId, nothing more to do
+            if (!runtimeId)
+                return;
+        }
+        else {
+            // runtimeId didn't change, skip
+            return;
+        }
+        const store = useConversationStore.getState();
+        // Check if we need to fetch (not already fetched and not currently fetching)
+        if (!store.needsFetch(runtimeId)) {
+            // Already fetched - load from in-memory store
+            const storedMessages = store.getMessages(runtimeId);
+            if (storedMessages.length > 0) {
+                setDisplayItems(storedMessages);
+            }
+            return;
+        }
+        // Mark as fetching to prevent duplicate requests
+        store.setFetching(runtimeId, true);
+        // Build the history endpoint URL
+        const endpoint = historyEndpoint ||
+            (protocol?.endpoint ? `${protocol.endpoint}/history` : null);
+        if (!endpoint) {
+            console.warn('[ChatBase] No history endpoint available for runtimeId:', runtimeId);
+            store.markFetched(runtimeId);
+            return;
+        }
+        // Fetch conversation history from server
+        const fetchHistory = async () => {
+            try {
+                const authToken = historyAuthToken || protocol?.authToken;
+                const headers = {
+                    'Content-Type': 'application/json',
+                };
+                if (authToken) {
+                    headers['Authorization'] = `Bearer ${authToken}`;
+                }
+                const response = await fetch(endpoint, {
+                    method: 'GET',
+                    headers,
+                    credentials: 'include',
+                });
+                if (!response.ok) {
+                    throw new Error(`Failed to fetch history: ${response.status} ${response.statusText}`);
+                }
+                const data = await response.json();
+                const messages = data.messages || [];
+                // Store in memory and update display
+                if (messages.length > 0) {
+                    store.setMessages(runtimeId, messages);
+                    setDisplayItems(messages);
+                }
+                store.markFetched(runtimeId);
+            }
+            catch (err) {
+                console.error('[ChatBase] Failed to fetch conversation history:', err);
+                store.markFetched(runtimeId);
+            }
+        };
+        fetchHistory();
+    }, [
+        runtimeId,
+        historyEndpoint,
+        historyAuthToken,
+        protocol?.endpoint,
+        protocol?.authToken,
+    ]);
+    // Keep in-memory store in sync with displayItems (for session persistence)
+    useEffect(() => {
+        if (runtimeId && displayItems.length > 0) {
+            // Filter to only save ChatMessage items (not tool call items)
+            const messagesToSave = displayItems.filter((item) => !isToolCallMessage(item));
+            if (messagesToSave.length > 0) {
+                useConversationStore.getState().setMessages(runtimeId, messagesToSave);
+            }
+        }
+    }, [runtimeId, displayItems]);
     // Derived state
     const messages = displayItems.filter((item) => !isToolCallMessage(item));
     const ready = true;
@@ -954,9 +1167,13 @@ onAuthorizationRequired, connectedIdentities, }) {
         if (useStoreMode) {
             clearStoreMessages();
         }
+        // Clear from conversation store if runtimeId is provided
+        if (runtimeId) {
+            useConversationStore.getState().clearMessages(runtimeId);
+        }
         onNewChat?.();
         headerButtons?.onNewChat?.();
-    }, [clearStoreMessages, onNewChat, headerButtons, useStoreMode]);
+    }, [clearStoreMessages, onNewChat, headerButtons, useStoreMode, runtimeId]);
     // Handle clear
     const handleClear = useCallback(() => {
         if (window.confirm('Clear all messages?')) {
@@ -965,10 +1182,14 @@ onAuthorizationRequired, connectedIdentities, }) {
             if (useStoreMode) {
                 clearStoreMessages();
             }
+            // Clear from conversation store if runtimeId is provided
+            if (runtimeId) {
+                useConversationStore.getState().clearMessages(runtimeId);
+            }
             onClear?.();
             headerButtons?.onClear?.();
         }
-    }, [clearStoreMessages, onClear, headerButtons, useStoreMode]);
+    }, [clearStoreMessages, onClear, headerButtons, useStoreMode, runtimeId]);
     // Not ready yet (store mode only)
     if (!ready) {
         return (_jsx(Box, { className: className, sx: {
@@ -1000,6 +1221,34 @@ onAuthorizationRequired, connectedIdentities, }) {
                     p: padding,
                 }, children: [_jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 2 }, children: [brandIcon || _jsx(AiAgentIcon, { colored: true, size: 20 }), title && (_jsx(Heading, { as: "h3", sx: { fontSize: 2, fontWeight: 'semibold' }, children: title })), headerContent] }), _jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 1 }, children: [headerButtons?.showNewChat && (_jsx(IconButton, { icon: PlusIcon, "aria-label": "New chat", variant: "invisible", size: "small", onClick: handleNewChat })), headerButtons?.showClear && messages.length > 0 && (_jsx(IconButton, { icon: TrashIcon, "aria-label": "Clear messages", variant: "invisible", size: "small", onClick: handleClear })), headerButtons?.showSettings && (_jsx(IconButton, { icon: GearIcon, "aria-label": "Settings", variant: "invisible", size: "small", onClick: headerButtons.onSettings })), headerActions] })] }) }));
     };
+    // Render token usage bar between input and selectors
+    const renderTokenUsage = () => {
+        if (!showTokenUsage)
+            return null;
+        // Show bar when we have any context data (totalTokens > 0 means agent is active)
+        const hasContext = agentUsage && !agentUsage.error && agentUsage.totalTokens > 0;
+        const hasTurn = agentUsage?.turnUsage &&
+            (agentUsage.turnUsage.inputTokens > 0 ||
+                agentUsage.turnUsage.outputTokens > 0);
+        const hasSession = agentUsage?.sessionUsage &&
+            (agentUsage.sessionUsage.inputTokens > 0 ||
+                agentUsage.sessionUsage.outputTokens > 0);
+        if (!hasContext)
+            return null;
+        return (_jsxs(Box, { sx: {
+                display: 'flex',
+                alignItems: 'center',
+                justifyContent: 'flex-start',
+                gap: 2,
+                py: 1,
+                px: padding,
+                bg: 'canvas.subtle',
+                flexWrap: 'nowrap',
+                overflow: 'hidden',
+                whiteSpace: 'nowrap',
+                minWidth: 0,
+            }, children: [_jsxs(Text, { sx: { fontSize: 0, color: 'fg.muted', flexShrink: 0 }, children: [_jsx(Text, { as: "span", sx: { fontWeight: 'semibold', color: 'fg.default', fontSize: 0 }, children: formatTokenCount(agentUsage.totalTokens) }), ' / ', formatTokenCount(agentUsage.contextWindow), ' ctx'] }), hasSession && (_jsxs(Text, { sx: { fontSize: 0, color: 'fg.muted', flexShrink: 0 }, children: ['· ', formatTokenCount(agentUsage.sessionUsage.inputTokens), _jsx(Text, { as: "span", sx: { color: 'success.fg', fontSize: 0 }, children: '▲' }), ' ', formatTokenCount(agentUsage.sessionUsage.outputTokens), _jsx(Text, { as: "span", sx: { color: 'attention.fg', fontSize: 0 }, children: '▼' })] })), hasTurn && (_jsxs(Text, { sx: { fontSize: 0, color: 'fg.muted', flexShrink: 0 }, children: ['· turn ', formatTokenCount(agentUsage.turnUsage.inputTokens), _jsx(Text, { as: "span", sx: { color: 'success.fg', fontSize: 0 }, children: '▲' }), ' ', formatTokenCount(agentUsage.turnUsage.outputTokens), _jsx(Text, { as: "span", sx: { color: 'attention.fg', fontSize: 0 }, children: '▼' })] }))] }));
+    };
     // Render empty state
     const renderEmptyState = () => {
         if (emptyState?.render) {
@@ -1263,7 +1512,9 @@ onAuthorizationRequired, connectedIdentities, }) {
                                         maxWidth: '85%',
                                         p: 2,
                                         borderRadius: 2,
-                                        backgroundColor: isUser ? 'accent.emphasis' : '#f6f8fa',
+                                        backgroundColor: isUser
+                                            ? 'accent.emphasis'
+                                            : 'canvas.subtle',
                                         color: isUser ? 'fg.onEmphasis' : 'fg.default',
                                         // Streamdown code block styling
                                         // Code block container
@@ -1279,10 +1530,10 @@ onAuthorizationRequired, connectedIdentities, }) {
                                             display: 'flex',
                                             alignItems: 'center',
                                             justifyContent: 'space-between',
-                                            backgroundColor: '#e1e4e8',
+                                            backgroundColor: 'canvas.subtle',
                                             padding: '8px 12px',
                                             fontSize: '12px',
-                                            color: '#586069',
+                                            color: 'fg.muted',
                                         },
                                         // Style the buttons in the header
                                         '& [data-streamdown="code-block-header"] button': {
@@ -1290,16 +1541,16 @@ onAuthorizationRequired, connectedIdentities, }) {
                                             border: 'none',
                                             cursor: 'pointer',
                                             padding: '4px',
-                                            color: '#586069',
+                                            color: 'fg.muted',
                                             borderRadius: '4px',
                                             '&:hover': {
-                                                backgroundColor: 'rgba(0,0,0,0.1)',
-                                                color: '#24292e',
+                                                backgroundColor: 'neutral.muted',
+                                                color: 'fg.default',
                                             },
                                         },
                                         // Code block body
                                         '& [data-streamdown="code-block-body"]': {
-                                            backgroundColor: '#f6f8fa',
+                                            backgroundColor: 'canvas.subtle',
                                             padding: '12px',
                                             margin: 0,
                                             overflow: 'auto',
@@ -1429,7 +1680,7 @@ onAuthorizationRequired, connectedIdentities, }) {
     // Render protocol mode input
     const renderProtocolInput = () => {
         const availableTools = configQuery.data?.builtinTools || [];
-        const models = configQuery.data?.models || [];
+        const models = availableModels || configQuery.data?.models || [];
         return (_jsxs(Box, { children: [_jsx(Box, { sx: {
                         p: padding,
                         borderTop: '1px solid',
@@ -1445,7 +1696,7 @@ onAuthorizationRequired, connectedIdentities, }) {
                                     maxHeight: '120px',
                                     overflow: 'hidden',
                                     transition: 'height 0.1s ease-out',
-                                }, rows: 1 }), isLoading ? (_jsx(IconButton, { icon: SquareCircleIcon, "aria-label": "Stop", onClick: handleStop, sx: { alignSelf: 'flex-end' } })) : (_jsx(IconButton, { icon: PaperAirplaneIcon, "aria-label": "Send", onClick: handleSend, disabled: !input.trim(), sx: { alignSelf: 'flex-end' } }))] }) }), (showModelSelector || showToolsMenu || showSkillsMenu) &&
+                                }, rows: 1 }), isLoading ? (_jsx(IconButton, { icon: SquareCircleIcon, "aria-label": "Stop", onClick: handleStop, sx: { alignSelf: 'flex-end' } })) : (_jsx(IconButton, { icon: PaperAirplaneIcon, "aria-label": "Send", onClick: handleSend, disabled: !input.trim(), sx: { alignSelf: 'flex-end' } }))] }) }), renderTokenUsage(), (showModelSelector || showToolsMenu || showSkillsMenu) &&
                     (configQuery.data || skillsQuery.data) && (_jsxs(Box, { sx: {
                         display: 'flex',
                         gap: 2,
@@ -1460,7 +1711,12 @@ onAuthorizationRequired, connectedIdentities, }) {
                                             maxHeight: '60vh',
                                             overflowY: 'auto',
                                         }, children: _jsxs(ActionList, { children: [codemodeEnabled && (_jsx(ActionList.Group, { title: "Codemode", children: _jsx(ActionList.Item, { disabled: true, children: _jsx(Text, { sx: { fontSize: 0, color: 'fg.muted' }, children: "MCP tools are accessible via Codemode meta-tools (search_tools, list_tool_names, execute_code)." }) }) })), configQuery.data?.mcpServers &&
-                                                    configQuery.data.mcpServers.length > 0 ? (configQuery.data.mcpServers.map(server => {
+                                                    configQuery.data.mcpServers.length > 0 ? (
+                                                // Filter to only show selected servers (if mcpServers was provided)
+                                                // If mcpServers is empty array, show no servers
+                                                configQuery.data.mcpServers.filter(server => !mcpServers || isServerSelected(server)).length > 0 ? (configQuery.data.mcpServers
+                                                    .filter(server => !mcpServers || isServerSelected(server))
+                                                    .map(server => {
                                                     const serverTools = enabledMcpTools.get(server.id);
                                                     const allToolNames = server.tools.map(t => t.name);
                                                     const enabledCount = serverTools?.size ?? 0;
@@ -1506,7 +1762,10 @@ onAuthorizationRequired, connectedIdentities, }) {
                                                                         color: 'fg.muted',
                                                                         fontStyle: 'italic',
                                                                     }, children: "Server unavailable" }) }))] }, server.id));
-                                                })) : (_jsx(ActionList.Group, { title: "Available Tools", children: availableTools.length > 0 ? (availableTools.map(tool => (_jsxs(ActionList.Item, { disabled: true, children: [_jsx(ActionList.LeadingVisual, { children: _jsx(Box, { sx: {
+                                                })) : (_jsx(ActionList.Group, { title: "MCP Servers", children: _jsx(ActionList.Item, { disabled: true, children: _jsx(Text, { sx: {
+                                                                color: 'fg.muted',
+                                                                fontStyle: 'italic',
+                                                            }, children: "No MCP servers selected" }) }) }))) : (_jsx(ActionList.Group, { title: "Available Tools", children: availableTools.length > 0 ? (availableTools.map(tool => (_jsxs(ActionList.Item, { disabled: true, children: [_jsx(ActionList.LeadingVisual, { children: _jsx(Box, { sx: {
                                                                         width: 8,
                                                                         height: 8,
                                                                         borderRadius: '50%',

package/lib/components/chat/components/index.d.ts

@@ -7,11 +7,13 @@ export { ChatMessages, type ChatMessagesProps } from './elements/ChatMessages';
 export { ChatInputPrompt, type ChatInputPromptProps, } from './elements/ChatInputPrompt';
 export { ChatSidebar, type ChatSidebarProps } from './ChatSidebar';
 export { ChatStandalone, type ChatStandaloneProps, type MessageHandler, } from './ChatStandalone';
-export { ChatBase, type ChatBaseProps, type ProtocolConfig, } from './base/ChatBase';
+export { ChatBase, type ChatBaseProps, type ProtocolConfig, type AgentRuntimeConfig, type AvatarConfig, type EmptyStateConfig, type HeaderButtonsConfig, type StreamingMessageOptions, } from './base/ChatBase';
 export { AgentDetails, type AgentDetailsProps } from './AgentDetails';
 export { AgentIdentity, IdentityCard, getTokenStatus, formatDuration, formatExpirationStatus, type AgentIdentityProps, type IdentityCardProps, type TokenStatus, } from './AgentIdentity';
 export { ContextUsage, type ContextUsageProps, type ContextDetailsResponse, } from './ContextUsage';
 export { ContextDistribution, type ContextDistributionProps, type ContextSnapshotResponse, } from './ContextDistribution';
+export { ContextPanel, type ContextPanelProps } from './ContextPanel';
+export { ContextInspector, type ContextInspectorProps, type FullContextResponse, } from './ContextInspector';
 export { ToolApprovalDialog, useToolApprovalDialog, type ToolApprovalDialogProps, } from './elements/ToolApprovalDialog';
 export { PoweredByTag, type PoweredByTagProps } from './elements/PoweredByTag';
 export { FloatingBrandButton, type FloatingBrandButtonProps, } from './elements/FloatingBrandButton';

package/lib/components/chat/components/index.js

@@ -16,6 +16,8 @@ export { AgentDetails } from './AgentDetails';
 export { AgentIdentity, IdentityCard, getTokenStatus, formatDuration, formatExpirationStatus, } from './AgentIdentity';
 export { ContextUsage, } from './ContextUsage';
 export { ContextDistribution, } from './ContextDistribution';
+export { ContextPanel } from './ContextPanel';
+export { ContextInspector, } from './ContextInspector';
 export { ToolApprovalDialog, useToolApprovalDialog, } from './elements/ToolApprovalDialog';
 export { PoweredByTag } from './elements/PoweredByTag';
 export { FloatingBrandButton, } from './elements/FloatingBrandButton';

package/lib/components/chat/extensions/ExtensionRegistry.d.ts

@@ -6,7 +6,7 @@
  */
 import type { ChatExtension, MessageRendererExtension, ActivityRendererExtension, ToolUIExtension, ProtocolEventExtension, PanelExtension } from '../types/extension';
 /** Internal extension type for registry organization */
-type InternalExtensionType = 'message-renderer' | 'activity-renderer' | 'tool-ui' | 'protocol-event' | 'panel';
+export type InternalExtensionType = 'message-renderer' | 'activity-renderer' | 'tool-ui' | 'protocol-event' | 'panel';
 /**
  * Extension registry class
  */

package/lib/components/chat/extensions/index.d.ts

@@ -4,6 +4,7 @@
  * @module components/chat/extensions
  */
 export { ExtensionRegistry, createMessageRenderer, createActivityRenderer, } from './ExtensionRegistry';
+export type { InternalExtensionType } from './ExtensionRegistry';
 export { createA2UIRenderer, A2UIExtensionImpl, type A2UIMessage, } from './A2UIExtension';
 export { createMCPUIRenderer, MCPUIExtensionImpl, type MCPUIMessage, type MCPUIResource, } from './MCPUIExtension';
 export type { ChatExtension, ExtensionType, MessageRendererExtension, ActivityRendererExtension, ToolUIExtension, ProtocolEventExtension, PanelExtension, A2UIExtension, } from '../types/extension';

package/lib/components/chat/index.d.ts

@@ -49,13 +49,13 @@
  * ```
  */
 export * from './types';
-export { useChatStore, useChatMessages, useChatLoading, useChatStreaming, useChatError, useChatTools, useChatOpen, useChatConfig, useChatReady, useChatInferenceProvider, useChatExtensionRegistry, defaultChatConfig, type ChatStore, type ChatConfig, } from './store';
+export { useChatStore, useChatMessages, useChatLoading, useChatStreaming, useChatError, useChatTools, useChatOpen, useChatConfig, useChatReady, useChatInferenceProvider, useChatExtensionRegistry, defaultChatConfig, type ChatStore, type ChatState, type ChatActions, type ChatConfig, type ToolCallState, } from './store';
 export { useChat, useFrontendTool, useBackendTool, ActionRegistrar, type UseChatReturn, type UseFrontendToolFn, } from '../../hooks';
 export { BaseInferenceProvider, DatalayerInferenceProvider, SelfHostedInferenceProvider, type DatalayerInferenceConfig, type SelfHostedInferenceConfig, } from './inference';
 export { BaseProtocolAdapter, AGUIAdapter, A2AAdapter, ACPAdapter, type AGUIAdapterConfig, type A2AAdapterConfig, type ACPAdapterConfig, type ACPSession, type ACPAgent, type ACPPendingPermission, } from './protocols';
 export { ToolExecutor, type ToolExecutionContext } from './tools';
 export { MiddlewarePipeline, createMiddleware, loggingMiddleware, createHITLMiddleware, type RequestContext, type ResponseContext, } from './middleware';
-export { ExtensionRegistry, createMessageRenderer, createActivityRenderer, createA2UIRenderer, A2UIExtensionImpl, type A2UIMessage, } from './extensions';
-export { ChatMessages, ChatInputPrompt, ChatSidebar, ChatStandalone, ChatBase, ToolApprovalDialog, useToolApprovalDialog, PoweredByTag, FloatingBrandButton, ChatHeader, MessagePart, TextPart, ReasoningPart, ToolPart, DynamicToolPart, ToolCallDisplay, Chat, ChatFloating, AgentDetails, AgentIdentity, IdentityCard, getTokenStatus, formatDuration, formatExpirationStatus, type ChatMessagesProps, type ChatInputPromptProps, type ChatSidebarProps, type ChatStandaloneProps, type MessageHandler, type ChatBaseProps, type ProtocolConfig, type ToolApprovalDialogProps, type PoweredByTagProps, type FloatingBrandButtonProps, type ChatFloatingProps, type ToolCallRenderContext, type ToolCallStatus, type RenderToolResult, type RespondCallback, type Suggestion, type RemoteConfig, type ModelConfig, type BuiltinTool, type MCPServerConfig, type MCPServerTool, type AgentDetailsProps, type AgentIdentityProps, type IdentityCardProps, type TokenStatus, type ChatHeaderProps, type ConnectionState, type MessagePartProps, type TextPartProps, type ReasoningPartProps, type ToolPartProps, type DynamicToolPartProps, type ToolCallDisplayProps, type ChatProps, type Transport, type Extension, } from './components';
+export { ExtensionRegistry, createMessageRenderer, createActivityRenderer, createA2UIRenderer, A2UIExtensionImpl, type A2UIMessage, type InternalExtensionType, } from './extensions';
+export { ChatMessages, ChatInputPrompt, ChatSidebar, ChatStandalone, ChatBase, ToolApprovalDialog, useToolApprovalDialog, PoweredByTag, FloatingBrandButton, ChatHeader, MessagePart, TextPart, ReasoningPart, ToolPart, DynamicToolPart, ToolCallDisplay, Chat, ChatFloating, AgentDetails, AgentIdentity, IdentityCard, getTokenStatus, formatDuration, formatExpirationStatus, type ChatMessagesProps, type ChatInputPromptProps, type ChatSidebarProps, type ChatStandaloneProps, type MessageHandler, type ChatBaseProps, type ProtocolConfig, type AgentRuntimeConfig, type AvatarConfig, type EmptyStateConfig, type HeaderButtonsConfig, type StreamingMessageOptions, type ToolApprovalDialogProps, type PoweredByTagProps, type FloatingBrandButtonProps, type ChatFloatingProps, type ToolCallRenderContext, type ToolCallStatus, type RenderToolResult, type RespondCallback, type Suggestion, type RemoteConfig, type ModelConfig, type BuiltinTool, type MCPServerConfig, type MCPServerTool, type AgentDetailsProps, type AgentIdentityProps, type IdentityCardProps, type TokenStatus, type ChatHeaderProps, type ConnectionState, type MessagePartProps, type TextPartProps, type ReasoningPartProps, type ToolPartProps, type DynamicToolPartProps, type ToolCallDisplayProps, type ChatProps, type Transport, type Extension, } from './components';
 export { requestAPI } from './handler';
 export { useKeyboardShortcuts, useChatKeyboardShortcuts, getShortcutDisplay, type KeyboardShortcut, type UseKeyboardShortcutsOptions, } from '../../hooks';