@datalayer/agent-runtimes 0.0.4 → 0.0.7

package/README.md

@@ -2,37 +2,34 @@
   ~ Copyright (c) 2025-2026 Datalayer, Inc.
   ~
   ~ BSD 3-Clause License
+-->
 
+[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io)
 
-So your goal is to require approval for all or only some tools on that MCP server? And then to present the approval to the user using the Vercel AI integration?
-yes, that is my goal to implement https://github.com/datalayer/jupyter-ai-agents/issues/53
+[![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer)
 
+# 🤖 Agent Runtimes
 
-As you know, Vercel AI Element are just React UI components, it is up to the developer to implement such features.
-I believe it's a little bit more than that because of the Data Stream Protocol they implement: https://ai-sdk.dev/docs/ai-sdk-ui/stream-protocol#data-stream-protocol, that we support as well: https://ai.pydantic.dev/ui/vercel-ai/, so you don't just get the components but also the automatic rendering of the chat interface based on the events that are coming from the server. Right?
-The vercel ai-element tool approval has been merged 3 weeks ago https://github.com/vercel/ai-elements/pull/163
-That's great! It doesn't look like they have events for that in the Data Stream Protocol yet though, so it's not clear how they expect agent backends to represent tool calls that require approval, or how they'll send approval to the backend...
-I guess having a kind of reference implementation for common use cases/patterns will be good for the community. I would say that kind of case (external mcp server with tool approval like done by eg claude desktop) is quite common, or will become common
-Definitely, which is partly why github.com/pydantic/pydantic-ai/issues/3295 exists, but there's currently a missing link in how approvals should be represented on the event stream going from agent to frontend, and in approval results going from frontend to agent. I'd expect Vercel AI to come up with something sooner rather than later since they have their own (TypeScript) agent SDK as well, which they'd then document in the Data Stream Protocol, which we can then support as well.
-In the meantime, you could do something like this:
-Wrap your MCPServer in ApprovalRequiredToolset (assuming all tools require approval). This means that when any of the tools are called, the agent run will end with a DeferredToolRequests object as result.output
-When you use VercelAIAdapter.dispatch_request, you can add an on_complete handler that is passed the AgentRunResult and can yield additional Vercel AI events. In that handler, you can check result.output , and if it's DeferredToolRequests , you can yield a custom pydantic_ai.ui.vercel_ai.response_types.DataChunk wrapping those deferred tool requests, which will be sent to frontend
-On the frontend you can handle that custom data part by showing the user the tool call that needs approval
-When the user approves, you'll need to perform a new request to the backend containing that approval, by listing the approved tool_call_ids on the SubmitMessage payload (I'm not an expert on how this part of AI Elements works...)
-In the backend endpoint, you then have to parse this value out of the SubmitMessage payload (similar to this: https://github.com/pydantic/ai-chat-ui/blob/b0ac18d2dbbef3cf1636c17a29fb7578a59b32c8/agent/chatbot/server.py#L115-L122), turn it into a DeferredToolResults object and pass it to the dispatch/run method as deferred_tool_results
+[![Github Actions Status](https://github.com/datalayer/agent-runtimes/actions/workflows/build.yml/badge.svg)](https://github.com/datalayer/agent-runtimes/actions/workflows/build.yml)
+[![PyPI - Version](https://img.shields.io/pypi/v/agent-runtimes)](https://pypi.org/project/agent-runtimes)
 
--->
+**Agent Runtimes** is a unified platform for deploying, managing, and interacting with AI agents across multiple protocols and frameworks. It provides both a Python server for hosting agents and React components for seamless integration into web and desktop applications.
 
-[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io)
+## 🎯 What is Agent Runtimes?
 
-[![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer)
+Agent Runtimes solves the complexity of deploying AI agents by providing:
 
-# 🤖 Agent Runtimes
+1. **Protocol Abstraction**: One agent, multiple protocols - deploy your agent once and access it through ACP, Vercel AI SDK, AG-UI, MCP-UI, or A2A without changing your code.
+
+2. **Framework Flexibility**: Write agents using your preferred framework (Pydantic AI, LangChain, Jupyter AI) while maintaining a consistent API.
 
-[![Github Actions Status](https://github.com/datalayer/code-sandboxes/workflows/Build/badge.svg)](https://github.com/datalayer/code-sandboxes/actions/workflows/build.yml)
-[![PyPI - Version](https://img.shields.io/pypi/v/code-sandboxes)](https://pypi.org/project/code-sandboxes)
+3. **Cloud Runtime Management**: Built-in integration with Datalayer Cloud Runtimes for launching and managing compute resources with Zustand-based state management.
 
-**Agent Runtimes** is a flexible server framework for exposing AI agents through multiple protocols.
+4. **UI Components**: Pre-built React components (ChatBase, ChatSidebar, ChatFloating) that connect to agents and execute tools directly in the browser.
+
+5. **Tool Ecosystem**: Seamless integration with MCP (Model Context Protocol) tools, custom tools, and built-in utilities for Jupyter notebooks and Lexical documents.
+
+![Agent Runtimes](https://images.datalayer.io/product/agent-runtimes/agent-runtimes-example-1.gif)
 
 ## 🌟 Features
 
@@ -49,8 +46,130 @@ In the backend endpoint, you then have to parse this value out of the SubmitMess
 - **Jupyter AI**: Notebook integration (adapter ready)
 
 ### Built-in Features
-- 🔌 **Adapter Architecture**: Easy to add new agents and protocols
+- 🔌 **Flexible Architecture**: Easy to add new agents and protocols
 - 🛠️ **Tool Support**: MCP, custom tools, built-in utilities
 - 📊 **Observability**: OpenTelemetry integration
 - 💾 **Persistence**: DBOS support for durable execution
 - 🔒 **Context Optimization**: LLM context management
+
+## 🏗️ Architecture
+
+Agent Runtimes consists of three main components:
+
+### 1. Python Server (`agent_runtimes/`)
+The backend server that hosts AI agents and handles protocol routing:
+- **Agent Adapters**: Unified interface for Pydantic AI, LangChain, and Jupyter AI
+- **Protocol Adapters**: Convert between different agent protocols (ACP, AG-UI, Vercel AI, etc.)
+- **FastAPI Server**: High-performance async server with automatic API documentation
+- **Tool Registry**: Manages and executes tools from various sources (MCP, custom, built-in)
+
+### 2. React Components (`src/components/`)
+Pre-built UI components for interacting with agents:
+- **ChatBase**: Core chat interface with customizable styling
+- **ChatSidebar**: Collapsible sidebar for agent interactions
+- **ChatFloating**: Floating chat widget for non-intrusive agent access
+- **All components support**: Frontend tool execution, markdown rendering, code highlighting, and real-time streaming
+
+### 3. Runtime Management (`src/runtime/`)
+Cloud runtime lifecycle management with Zustand store:
+- **Launch & Connect**: Create new cloud runtimes or connect to existing ones
+- **Agent Creation**: Automatically create and configure agents on runtimes
+- **State Management**: Track runtime status, agent connections, and errors
+- **Hooks**: React hooks for easy integration (`useAgentRuntime`, `useRuntimeStore`)
+
+## 🚀 Use Cases
+
+### Notebook AI Assistant
+Add an AI agent to Jupyter notebooks that can:
+- Execute cells, insert code, and modify notebook content
+- Explain code and data analysis
+- Debug errors and suggest improvements
+
+```tsx
+import { NotebookAgentsRuntime } from '@datalayer/agent-runtimes';
+
+<NotebookAgentsRuntime
+  notebookId={notebookId}
+  environmentName="python-simple"
+  runtimeName={runtimeName}
+  serviceManager={serviceManager}
+/>
+```
+
+### Document Editor AI
+Integrate AI into Lexical-based document editors:
+- Insert headings, lists, code blocks, and formatted text
+- Summarize content and proofread text
+- Generate ideas and help with writing
+
+```tsx
+import { DocumentAgentRuntime } from '@datalayer/agent-runtimes';
+
+<DocumentAgentRuntime
+  documentId={documentId}
+  environmentName="python-simple"
+  runtimeName={runtimeName}
+  serviceManager={serviceManager}
+/>
+```
+
+### Custom Agent Deployment
+Deploy your own Pydantic AI agent with custom tools:
+
+```python
+from agent_runtimes import AgentRuntimesApp
+from pydantic_ai import Agent
+
+# Create your agent
+agent = Agent(
+    model='anthropic:claude-sonnet-4-5',
+    system_prompt='You are a helpful assistant.',
+)
+
+# Launch the server
+app = AgentRuntimesApp()
+app.add_agent(agent, name='my-agent', transport='ag-ui')
+app.run(port=8000)
+```
+
+## 🔧 Key Concepts
+
+### Protocols
+Agent Runtimes supports multiple protocols for agent communication:
+
+- **AG-UI**: Lightweight protocol for web UIs (POST-based, Pydantic AI native)
+- **ACP**: WebSocket-based Agent Client Protocol for real-time interaction
+- **Vercel AI SDK**: Compatible with Vercel's AI SDK streaming
+- **MCP-UI**: Model Context Protocol with UI resources
+- **A2A**: Agent-to-agent communication protocol
+
+### Tools
+Tools extend agent capabilities by allowing them to perform actions:
+
+- **Frontend Tools**: Execute in the browser (notebook editing, document manipulation)
+- **MCP Tools**: Tools from Model Context Protocol servers
+- **Custom Tools**: Your own Python functions decorated with tool metadata
+- **Built-in Tools**: File operations, web search, code execution
+- **Code Mode**: Tool discovery includes `output_schema` and `input_examples` for reliable calls; code execution returns `stdout`/`stderr` and a summarized `result`.
+
+### Runtime Management
+Cloud runtimes provide compute resources for agents:
+
+```typescript
+import { useAgentRuntime } from '@datalayer/agent-runtimes/lib/runtime';
+
+const { isReady, endpoint, tools, launchRuntime } = useAgentRuntime({
+  autoCreateAgent: true,
+  agentConfig: {
+    model: 'anthropic:claude-sonnet-4-5',
+    systemPrompt: 'You are a helpful AI assistant.',
+  },
+});
+
+// Launch a new runtime
+await launchRuntime({
+  environmentName: 'python-simple',
+  creditsLimit: 100,
+  type: 'notebook',
+});
+```

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

@@ -7,7 +7,7 @@ export interface AgentDetailsProps {
     url: string;
     /** Number of messages in conversation */
     messageCount: number;
-    /** Agent ID */
+    /** Agent ID for context usage tracking */
     agentId?: string;
     /** Callback to go back to chat view */
     onBack: () => void;

package/lib/components/chat/components/AgentDetails.js

@@ -5,53 +5,12 @@ import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
  * AgentDetails component - Shows detailed information about the agent
  * including name, protocol, URL, message count, and context details.
  */
-import { ArrowLeftIcon, GlobeIcon, CommentDiscussionIcon, DatabaseIcon, FileIcon, ToolsIcon, ClockIcon, CheckCircleIcon, XCircleIcon, } from '@primer/octicons-react';
-import { Box, Button, Heading, IconButton, Text, Label, ProgressBar, Spinner, } from '@primer/react';
+import { ArrowLeftIcon, GlobeIcon, CommentDiscussionIcon, CheckCircleIcon, XCircleIcon, } from '@primer/octicons-react';
+import { Box, Button, Heading, IconButton, Text, Label, Spinner, } from '@primer/react';
 import { AiAgentIcon } from '@datalayer/icons-react';
 import { useQuery } from '@tanstack/react-query';
-// Mock context data for display
-const MOCK_CONTEXT_DATA = {
-    name: 'Context',
-    totalTokens: 2520000,
-    usedTokens: 1523552,
-    children: [
-        {
-            name: 'Files',
-            value: 450000,
-            children: [
-                { name: 'app.py', value: 125000 },
-                { name: 'models.py', value: 98000 },
-                { name: 'routes.py', value: 112000 },
-                { name: 'utils.py', value: 115000 },
-            ],
-        },
-        {
-            name: 'Messages',
-            value: 380000,
-            children: [
-                { name: 'User messages', value: 180000 },
-                { name: 'Assistant responses', value: 200000 },
-            ],
-        },
-        {
-            name: 'Tools',
-            value: 220000,
-            children: [
-                { name: 'Code execution', value: 95000 },
-                { name: 'File operations', value: 75000 },
-                { name: 'Search', value: 50000 },
-            ],
-        },
-        {
-            name: 'Memory',
-            value: 473552,
-            children: [
-                { name: 'Short term', value: 150000 },
-                { name: 'Long term', value: 323552 },
-            ],
-        },
-    ],
-};
+import { ContextUsage } from './ContextUsage';
+import { ContextDistribution } from './ContextDistribution';
 function getLocalApiBase() {
     if (typeof window === 'undefined') {
         return '';
@@ -61,40 +20,10 @@ function getLocalApiBase() {
         ? 'http://127.0.0.1:8765'
         : '';
 }
-/**
- * Format token count for display
- */
-function formatTokens(tokens) {
-    if (tokens >= 1000000) {
-        return `${(tokens / 1000000).toFixed(1)}M`;
-    }
-    if (tokens >= 1000) {
-        return `${(tokens / 1000).toFixed(1)}K`;
-    }
-    return tokens.toString();
-}
-/**
- * Get icon for context category
- */
-function getCategoryIcon(name) {
-    switch (name.toLowerCase()) {
-        case 'files':
-            return FileIcon;
-        case 'messages':
-            return CommentDiscussionIcon;
-        case 'tools':
-            return ToolsIcon;
-        case 'memory':
-            return DatabaseIcon;
-        default:
-            return ClockIcon;
-    }
-}
 /**
  * AgentDetails component displays comprehensive information about the agent.
  */
 export function AgentDetails({ name = 'AI Agent', protocol, url, messageCount, agentId, onBack, }) {
-    const contextUsagePercent = (MOCK_CONTEXT_DATA.usedTokens / MOCK_CONTEXT_DATA.totalTokens) * 100;
     // Fetch MCP toolsets status
     const { data: mcpStatus, isLoading: mcpLoading } = useQuery({
         queryKey: ['mcp-toolsets-status'],
@@ -203,31 +132,17 @@ export function AgentDetails({ name = 'AI Agent', protocol, url, messageCount, a
                                                                 pl: 4,
                                                                 whiteSpace: 'pre-wrap',
                                                                 wordBreak: 'break-word',
-                                                            }, children: error.split('\n')[0] })] }, server)))] }))] })) : (_jsx(Text, { sx: { fontSize: 1, color: 'fg.muted' }, children: "Failed to load MCP status" })) })] }), _jsxs(Box, { children: [_jsx(Heading, { as: "h4", sx: {
+                                                            }, children: error.split('\n')[0] })] }, server)))] }))] })) : (_jsx(Text, { sx: { fontSize: 1, color: 'fg.muted' }, children: "Failed to load MCP status" })) })] }), agentId && _jsx(ContextUsage, { agentId: agentId }), agentId && (_jsxs(Box, { sx: { mt: 3 }, children: [_jsx(Heading, { as: "h4", sx: {
                                     fontSize: 1,
                                     fontWeight: 'semibold',
                                     mb: 2,
                                     color: 'fg.muted',
-                                }, children: "Context Usage" }), _jsxs(Box, { sx: {
+                                }, children: "Current Context Distribution" }), _jsx(Box, { sx: {
                                     p: 3,
                                     bg: 'canvas.subtle',
                                     borderRadius: 2,
                                     border: '1px solid',
                                     borderColor: 'border.default',
-                                }, children: [_jsxs(Box, { sx: { mb: 3 }, children: [_jsxs(Box, { sx: {
-                                                    display: 'flex',
-                                                    justifyContent: 'space-between',
-                                                    mb: 1,
-                                                }, children: [_jsxs(Text, { sx: { fontSize: 1, fontWeight: 'semibold' }, children: [formatTokens(MOCK_CONTEXT_DATA.usedTokens), " /", ' ', formatTokens(MOCK_CONTEXT_DATA.totalTokens), " tokens"] }), _jsxs(Text, { sx: { fontSize: 1, color: 'fg.muted' }, children: [contextUsagePercent.toFixed(0), "%"] })] }), _jsx(ProgressBar, { progress: contextUsagePercent, sx: { height: 8 }, bg: contextUsagePercent > 80
-                                                    ? 'danger.emphasis'
-                                                    : 'accent.emphasis' })] }), _jsx(Box, { sx: { display: 'flex', flexDirection: 'column', gap: 2 }, children: MOCK_CONTEXT_DATA.children.map(category => {
-                                            const CategoryIcon = getCategoryIcon(category.name);
-                                            const categoryPercent = (category.value / MOCK_CONTEXT_DATA.totalTokens) * 100;
-                                            return (_jsxs(Box, { sx: {
-                                                    display: 'flex',
-                                                    alignItems: 'center',
-                                                    gap: 2,
-                                                }, children: [_jsx(Box, { sx: { color: 'fg.muted', width: 20 }, children: _jsx(CategoryIcon, { size: 16 }) }), _jsx(Text, { sx: { fontSize: 1, flex: 1 }, children: category.name }), _jsx(Text, { sx: { fontSize: 0, color: 'fg.muted', minWidth: 60 }, children: formatTokens(category.value) }), _jsx(Box, { sx: { width: 80 }, children: _jsx(ProgressBar, { progress: categoryPercent, sx: { height: 4 } }) })] }, category.name));
-                                        }) })] })] }), _jsx(Box, { sx: { mt: 2 }, children: _jsx(Button, { variant: "primary", onClick: onBack, sx: { width: '100%' }, children: "Back to Chat" }) })] })] }));
+                                }, children: _jsx(ContextDistribution, { agentId: agentId, height: "250px" }) })] })), _jsx(Box, { sx: { mt: 2 }, children: _jsx(Button, { variant: "primary", onClick: onBack, sx: { width: '100%' }, children: "Back to Chat" }) })] })] }));
 }
 export default AgentDetails;

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

@@ -49,6 +49,10 @@ export interface ChatProps {
     showModelSelector?: boolean;
     /** Show tools menu (fetched from /configure endpoint) */
     showToolsMenu?: boolean;
+    /** Initial model ID to select (e.g., 'openai:gpt-4o-mini') */
+    initialModel?: string;
+    /** Initial MCP server IDs to enable (others will be disabled) */
+    initialMcpServers?: string[];
     /** Clear messages when component mounts or agentId changes */
     clearOnMount?: boolean;
     /** Suggestions to show in empty state */
@@ -101,5 +105,5 @@ export interface ChatProps {
  * />
  * ```
  */
-export declare function Chat({ transport, extensions: _extensions, baseUrl, wsUrl, agentId, placeholder, title, autoConnect: _autoConnect, streaming: _streaming, onMessageSent: _onMessageSent, onMessageReceived: _onMessageReceived, onDisconnect, onLogout: _onLogout, onCollapsePanel: _onCollapsePanel, className, height, showHeader, showModelSelector, showToolsMenu, clearOnMount: _clearOnMount, suggestions, submitOnSuggestionClick, description, autoFocus, }: ChatProps): import("react/jsx-runtime").JSX.Element;
+export declare function Chat({ transport, extensions: _extensions, baseUrl, wsUrl, agentId, placeholder, title, autoConnect: _autoConnect, streaming: _streaming, onMessageSent: _onMessageSent, onMessageReceived: _onMessageReceived, onDisconnect, onLogout: _onLogout, onCollapsePanel: _onCollapsePanel, className, height, showHeader, showModelSelector, showToolsMenu, initialModel, initialMcpServers, clearOnMount: _clearOnMount, suggestions, submitOnSuggestionClick, description, autoFocus, }: ChatProps): import("react/jsx-runtime").JSX.Element;
 export default Chat;

package/lib/components/chat/components/Chat.js

@@ -129,11 +129,22 @@ function getProtocolType(transport) {
  * />
  * ```
  */
-export function Chat({ transport, extensions: _extensions, baseUrl = 'http://localhost:8765', wsUrl, agentId, placeholder = 'Type your message...', title, autoConnect: _autoConnect = true, streaming: _streaming = true, onMessageSent: _onMessageSent, onMessageReceived: _onMessageReceived, onDisconnect, onLogout: _onLogout, onCollapsePanel: _onCollapsePanel, className, height = '600px', showHeader = true, showModelSelector = true, showToolsMenu = true, clearOnMount: _clearOnMount = true, suggestions, submitOnSuggestionClick = true, description, autoFocus = false, }) {
+export function Chat({ transport, extensions: _extensions, baseUrl = 'http://localhost:8765', wsUrl, agentId, placeholder = 'Type your message...', title, autoConnect: _autoConnect = true, streaming: _streaming = true, onMessageSent: _onMessageSent, onMessageReceived: _onMessageReceived, onDisconnect, onLogout: _onLogout, onCollapsePanel: _onCollapsePanel, className, height = '600px', showHeader = true, showModelSelector = true, showToolsMenu = true, initialModel, initialMcpServers, clearOnMount: _clearOnMount = true, suggestions, submitOnSuggestionClick = true, description, autoFocus = false, }) {
     const [error, setError] = useState(null);
     const [isInitializing, setIsInitializing] = useState(true);
     const [showDetails, setShowDetails] = useState(false);
     const [messageCount, setMessageCount] = useState(0);
+    const [focusTrigger, setFocusTrigger] = useState(0);
+    // Focus the input when returning from details view
+    useEffect(() => {
+        if (!showDetails) {
+            // Small delay to ensure the chat view is visible before focusing
+            const timer = setTimeout(() => {
+                setFocusTrigger(prev => prev + 1);
+            }, 100);
+            return () => clearTimeout(timer);
+        }
+    }, [showDetails]);
     // Build protocol config based on transport
     const protocolConfig = useMemo(() => {
         try {
@@ -248,28 +259,26 @@ export function Chat({ transport, extensions: _extensions, baseUrl = 'http://loc
                 bg: 'canvas.default',
             }, children: [_jsx(Spinner, { size: "large" }), _jsxs(Text, { sx: { mt: 3, color: 'fg.muted' }, children: ["Connecting to ", transport.toUpperCase().replace('-', ' '), " agent..."] })] }));
     }
-    // Show agent details view
-    if (showDetails) {
-        return (_jsx(Box, { className: className, sx: {
-                position: 'relative',
-                height,
-                bg: 'canvas.default',
-                display: 'flex',
-                flexDirection: 'column',
-            }, children: _jsx(AgentDetails, { name: title || 'AI Agent', protocol: transport, url: protocolConfig?.endpoint || baseUrl, messageCount: messageCount, agentId: agentId, onBack: () => setShowDetails(false) }) }));
-    }
-    return (_jsx(Box, { className: className, sx: {
+    return (_jsxs(Box, { className: className, sx: {
             position: 'relative',
             height,
             bg: 'canvas.default',
             display: 'flex',
             flexDirection: 'column',
-        }, children: _jsx(ChatBase, { title: title, showHeader: showHeader, useStore: false, protocol: protocolConfig, placeholder: placeholder, description: description, suggestions: suggestions, submitOnSuggestionClick: submitOnSuggestionClick, autoFocus: autoFocus, headerContent: _jsx(IconButton, { icon: InfoIcon, "aria-label": "Agent details", variant: "invisible", size: "small", onClick: () => setShowDetails(true) }), showModelSelector: showModelSelector, showToolsMenu: showToolsMenu, onNewChat: handleNewChat, onMessagesChange: messages => setMessageCount(messages.length), headerButtons: {
-                showNewChat: true,
-                showClear: true,
-                onNewChat: handleNewChat,
-            }, avatarConfig: {
-                showAvatars: true,
-            }, backgroundColor: "canvas.default" }) }));
+        }, children: [_jsx(Box, { sx: {
+                    display: showDetails ? 'flex' : 'none',
+                    flexDirection: 'column',
+                    height: '100%',
+                }, children: _jsx(AgentDetails, { name: title || 'AI Agent', protocol: transport, url: protocolConfig?.endpoint || baseUrl, messageCount: messageCount, agentId: agentId, onBack: () => setShowDetails(false) }) }), _jsx(Box, { sx: {
+                    display: showDetails ? 'none' : 'flex',
+                    flexDirection: 'column',
+                    height: '100%',
+                }, children: _jsx(ChatBase, { title: title, showHeader: showHeader, protocol: protocolConfig, placeholder: placeholder, description: description, suggestions: suggestions, submitOnSuggestionClick: submitOnSuggestionClick, autoFocus: autoFocus, headerContent: _jsx(IconButton, { icon: InfoIcon, "aria-label": "Agent details", variant: "invisible", size: "small", onClick: () => setShowDetails(true) }), showModelSelector: showModelSelector, showToolsMenu: showToolsMenu, initialModel: initialModel, initialMcpServers: initialMcpServers, onNewChat: handleNewChat, onMessagesChange: messages => setMessageCount(messages.length), headerButtons: {
+                        showNewChat: true,
+                        showClear: true,
+                        onNewChat: handleNewChat,
+                    }, avatarConfig: {
+                        showAvatars: true,
+                    }, backgroundColor: "canvas.default", focusTrigger: focusTrigger }) })] }));
 }
 export default Chat;

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

@@ -0,0 +1,47 @@
+/**
+ * Distribution child item
+ */
+interface DistributionChild {
+    name: string;
+    value: number;
+    children?: DistributionChild[];
+}
+/**
+ * Distribution data for treemap
+ */
+interface Distribution {
+    name: string;
+    value: number;
+    children: DistributionChild[];
+}
+/**
+ * Context snapshot response from API
+ */
+export interface ContextSnapshotResponse {
+    agentId: string;
+    systemPrompts: string[];
+    systemPromptTokens: number;
+    messages: Array<{
+        role: string;
+        content: string;
+        estimatedTokens: number;
+        timestamp: string | null;
+    }>;
+    userMessageTokens: number;
+    assistantMessageTokens: number;
+    totalTokens: number;
+    contextWindow: number;
+    distribution: Distribution;
+    error?: string;
+}
+export interface ContextDistributionProps {
+    /** Agent ID for fetching context snapshot */
+    agentId: string;
+    /** Height of the chart */
+    height?: string;
+}
+/**
+ * ContextDistribution component displays context distribution as a treemap.
+ */
+export declare function ContextDistribution({ agentId, height, }: ContextDistributionProps): import("react/jsx-runtime").JSX.Element;
+export default ContextDistribution;

package/lib/components/chat/components/ContextDistribution.js

@@ -0,0 +1,146 @@
+import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
+// Copyright (c) 2025-2026 Datalayer, Inc.
+// Distributed under the terms of the Modified BSD License.
+/**
+ * ContextDistribution component - Shows context distribution as a treemap.
+ */
+import { Box, Text, Spinner, Button } from '@primer/react';
+import { ListUnorderedIcon } from '@primer/octicons-react';
+import { useQuery } from '@tanstack/react-query';
+import ReactECharts from 'echarts-for-react';
+import { useState } from 'react';
+function getLocalApiBase() {
+    if (typeof window === 'undefined') {
+        return '';
+    }
+    const host = window.location.hostname;
+    return host === 'localhost' || host === '127.0.0.1'
+        ? 'http://127.0.0.1:8765'
+        : '';
+}
+/**
+ * Format token count for display
+ */
+function formatTokens(tokens) {
+    if (tokens >= 1000000) {
+        return `${(tokens / 1000000).toFixed(1)}M`;
+    }
+    if (tokens >= 1000) {
+        return `${(tokens / 1000).toFixed(1)}K`;
+    }
+    return tokens.toString();
+}
+/**
+ * ContextDistribution component displays context distribution as a treemap.
+ */
+export function ContextDistribution({ agentId, height = '250px', }) {
+    const [showDetails, setShowDetails] = useState(false);
+    const { data: snapshotData, isLoading, error, } = useQuery({
+        queryKey: ['context-snapshot', agentId],
+        queryFn: async () => {
+            const apiBase = getLocalApiBase();
+            const response = await fetch(`${apiBase}/api/v1/configure/agents/${encodeURIComponent(agentId)}/context-snapshot`);
+            if (!response.ok) {
+                throw new Error('Failed to fetch context snapshot');
+            }
+            return response.json();
+        },
+        refetchInterval: 10000, // Refresh every 10 seconds
+        refetchOnMount: 'always',
+        staleTime: 0,
+    });
+    if (isLoading) {
+        return (_jsxs(Box, { sx: {
+                p: 3,
+                display: 'flex',
+                alignItems: 'center',
+                justifyContent: 'center',
+                height,
+            }, children: [_jsx(Spinner, { size: "small" }), _jsx(Text, { sx: { ml: 2, fontSize: 1, color: 'fg.muted' }, children: "Loading context distribution..." })] }));
+    }
+    if (error || !snapshotData) {
+        return (_jsx(Box, { sx: {
+                p: 3,
+                bg: 'canvas.subtle',
+                borderRadius: 2,
+                border: '1px solid',
+                borderColor: 'border.default',
+            }, children: _jsx(Text, { sx: { fontSize: 1, color: 'fg.muted' }, children: "Failed to load context distribution" }) }));
+    }
+    const { distribution, totalTokens, contextWindow } = snapshotData;
+    const hasData = distribution.children && distribution.children.length > 0;
+    // ECharts option for treemap
+    const chartOption = {
+        tooltip: {
+            formatter: (info) => {
+                return `${info.name}: ${formatTokens(info.value)} tokens`;
+            },
+        },
+        series: [
+            {
+                type: 'treemap',
+                data: hasData ? distribution.children : [{ name: 'No data', value: 1 }],
+                roam: false,
+                breadcrumb: {
+                    show: false,
+                },
+                label: {
+                    show: true,
+                    formatter: '{b}',
+                    fontSize: 12,
+                },
+                itemStyle: {
+                    borderColor: '#fff',
+                    borderWidth: 2,
+                },
+                levels: [
+                    {
+                        itemStyle: {
+                            borderColor: '#777',
+                            borderWidth: 0,
+                            gapWidth: 1,
+                        },
+                    },
+                    {
+                        itemStyle: {
+                            borderColor: '#555',
+                            borderWidth: 5,
+                            gapWidth: 1,
+                        },
+                        colorSaturation: [0.35, 0.5],
+                    },
+                    {
+                        colorSaturation: [0.35, 0.5],
+                    },
+                ],
+            },
+        ],
+    };
+    return (_jsxs(Box, { children: [_jsxs(Box, { sx: {
+                    display: 'flex',
+                    justifyContent: 'space-between',
+                    alignItems: 'center',
+                    mb: 2,
+                }, children: [_jsxs(Text, { sx: { fontSize: 1, fontWeight: 'bold' }, children: ["Current Context (", formatTokens(totalTokens), " /", ' ', formatTokens(contextWindow), " tokens)"] }), _jsx(Button, { size: "small", variant: "invisible", onClick: () => setShowDetails(!showDetails), leadingVisual: ListUnorderedIcon, children: showDetails ? 'Hide Details' : 'Show Details' })] }), hasData ? (_jsx(ReactECharts, { option: chartOption, style: { height }, opts: { renderer: 'svg' } })) : (_jsx(Box, { sx: {
+                    p: 4,
+                    bg: 'canvas.subtle',
+                    borderRadius: 2,
+                    textAlign: 'center',
+                    height,
+                    display: 'flex',
+                    alignItems: 'center',
+                    justifyContent: 'center',
+                }, children: _jsx(Text, { sx: { color: 'fg.muted', fontSize: 1 }, children: "No context data yet. Start a conversation to see context distribution." }) })), showDetails && hasData && (_jsxs(Box, { sx: {
+                    mt: 3,
+                    p: 2,
+                    border: '1px solid',
+                    borderColor: 'border.default',
+                    borderRadius: 2,
+                    bg: 'canvas.default',
+                    fontFamily: 'mono',
+                    fontSize: 0,
+                }, children: [_jsx(Text, { sx: { fontWeight: 'bold', display: 'block', mb: 2 }, children: "Context Breakdown:" }), snapshotData.systemPromptTokens > 0 && (_jsxs(Box, { sx: { mb: 2 }, children: [_jsxs(Text, { sx: { fontWeight: 'bold' }, children: ["System Prompts: ", formatTokens(snapshotData.systemPromptTokens), ' ', "tokens"] }), snapshotData.systemPrompts.map((prompt, idx) => (_jsxs(Text, { sx: { display: 'block', ml: 3, mt: 1, color: 'fg.muted' }, children: ["\u2022", ' ', prompt.length > 100 ? prompt.slice(0, 100) + '...' : prompt] }, idx)))] })), (snapshotData.userMessageTokens > 0 ||
+                        snapshotData.assistantMessageTokens > 0) && (_jsxs(Box, { sx: { mb: 2 }, children: [_jsxs(Text, { sx: { fontWeight: 'bold' }, children: ["Messages:", ' ', formatTokens(snapshotData.userMessageTokens +
+                                        snapshotData.assistantMessageTokens), ' ', "tokens"] }), _jsxs(Box, { sx: { ml: 3, mt: 1 }, children: [_jsxs(Text, { sx: { display: 'block' }, children: ["\u2022 User Messages:", ' ', formatTokens(snapshotData.userMessageTokens), " tokens"] }), _jsxs(Text, { sx: { display: 'block' }, children: ["\u2022 Assistant Responses:", ' ', formatTokens(snapshotData.assistantMessageTokens), " tokens"] })] })] }))] }))] }));
+}
+export default ContextDistribution;

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

@@ -0,0 +1,33 @@
+/**
+ * Context category child item
+ */
+interface ContextCategoryChild {
+    name: string;
+    value: number;
+}
+/**
+ * Context category with children
+ */
+interface ContextCategory {
+    name: string;
+    value: number;
+    children: ContextCategoryChild[];
+}
+/**
+ * Context details response from API
+ */
+export interface ContextDetailsResponse {
+    name: string;
+    totalTokens: number;
+    usedTokens: number;
+    children: ContextCategory[];
+}
+export interface ContextUsageProps {
+    /** Agent ID for fetching context details (required) */
+    agentId: string;
+}
+/**
+ * ContextUsage component displays token usage breakdown by category.
+ */
+export declare function ContextUsage({ agentId }: ContextUsageProps): import("react/jsx-runtime").JSX.Element;
+export default ContextUsage;