@datalayer/agent-runtimes 0.0.2 → 0.0.4

package/lib/components/chat/protocols/ACPAdapter.js

@@ -221,9 +221,14 @@ export class ACPAdapter extends BaseProtocolAdapter {
                     .filter(c => c.type === 'text')
                     .map(c => c.text || '')
                     .join('');
+            if (_options?.model) {
+                console.log('[ACPAdapter] Sending with model:', _options.model);
+            }
             await this.sendRequest(AGENT_METHODS.session_prompt, {
                 sessionId: this.session.sessionId,
                 content: [{ type: 'text', text: content }],
+                // Include model for per-request model override
+                ...(_options?.model && { metadata: { model: _options.model } }),
             });
         }
         catch (error) {
@@ -450,6 +455,14 @@ export class ACPAdapter extends BaseProtocolAdapter {
                 timestamp: new Date(),
             });
         }
+        // Handle error events - emit as error, not as message
+        if (params.error) {
+            this.emit({
+                type: 'error',
+                error: new Error(params.error),
+                timestamp: new Date(),
+            });
+        }
     }
     /**
      * Handle permission request from agent

package/lib/components/chat/protocols/AGUIAdapter.d.ts

@@ -52,6 +52,8 @@ export declare class AGUIAdapter extends BaseProtocolAdapter {
         metadata?: Record<string, unknown>;
         /** Full conversation history to send with the message */
         messages?: ChatMessage[];
+        /** Model to use for this request (overrides agent default) */
+        model?: string;
     }): Promise<void>;
     /**
      * Send tool result back through AG-UI and continue the conversation

package/lib/components/chat/protocols/AGUIAdapter.js

@@ -135,7 +135,12 @@ export class AGUIAdapter extends BaseProtocolAdapter {
             tools: options?.tools || [],
             context: [],
             forwardedProps: null,
+            // Include model for per-request model override
+            ...(options?.model && { model: options.model }),
         };
+        if (options?.model) {
+            console.log('[AGUIAdapter] Sending with model:', options.model);
+        }
         try {
             const response = await fetch(this.aguiConfig.baseUrl, {
                 method: 'POST',

package/lib/components/chat/protocols/BaseProtocolAdapter.d.ts

@@ -35,6 +35,8 @@ export declare abstract class BaseProtocolAdapter implements ProtocolAdapter {
         metadata?: Record<string, unknown>;
         /** Full conversation history to send with the message */
         messages?: ChatMessage[];
+        /** Model to use for this request (overrides agent default) */
+        model?: string;
     }): Promise<void>;
     /**
      * Send tool execution result back

package/lib/components/chat/protocols/VercelAIAdapter.d.ts

@@ -57,6 +57,12 @@ export declare class VercelAIAdapter extends BaseProtocolAdapter {
         tools?: ToolDefinition[];
         threadId?: string;
         metadata?: Record<string, unknown>;
+        /** Model to use for this request (overrides agent default) */
+        model?: string;
+        /** Full conversation history to send with the message */
+        messages?: ChatMessage[];
+        /** Builtin tools / MCP tools to enable for this request */
+        builtinTools?: string[];
     }): Promise<void>;
     /**
      * Parse SSE stream from Vercel AI

package/lib/components/chat/protocols/VercelAIAdapter.js

@@ -105,7 +105,17 @@ export class VercelAIAdapter extends BaseProtocolAdapter {
                 trigger: 'submit-message',
                 // Optional fields based on Pydantic AI's Vercel adapter
                 ...(options?.tools && { tools: options.tools }),
+                // Model override for per-request model selection
+                ...(options?.model && { model: options.model }),
+                // Builtin tools / MCP tools to enable
+                ...(options?.builtinTools &&
+                    options.builtinTools.length > 0 && {
+                    builtinTools: options.builtinTools,
+                }),
             };
+            if (options?.model) {
+                console.log('[VercelAIAdapter] Sending with model:', options.model);
+            }
             // Merge custom headers with defaults
             const headers = {
                 'Content-Type': 'application/json',

package/lib/components/chat/types/protocol.d.ts

@@ -115,6 +115,8 @@ export interface ProtocolAdapter {
         metadata?: Record<string, unknown>;
         /** Full conversation history to send with the message */
         messages?: ChatMessage[];
+        /** Model to use for this request (overrides agent default) */
+        model?: string;
     }): Promise<void>;
     /**
      * Send tool execution result back
@@ -161,6 +163,8 @@ export declare namespace AGUI {
             content: string;
         }>;
         forwardedProps: Record<string, unknown> | null;
+        /** Optional model override for per-request model selection */
+        model?: string;
     }
     interface Event {
         type: string;

package/lib/examples/AgentRuntimeCustomExample.d.ts

@@ -1,9 +1,10 @@
 /**
  * Agent Runtime Custom Example Component
  *
- * Demonstrates the unified Chat component with Simple transport
+ * Demonstrates the unified Chat component with AG-UI transport
  * and all necessary providers:
  * - QueryClientProvider for data fetching
+ * - Auto-creates agent on the server if it doesn't exist
  */
 declare const AgentRuntimeCustomExample: React.FC;
 export default AgentRuntimeCustomExample;

package/lib/examples/AgentRuntimeCustomExample.js

@@ -3,7 +3,8 @@ 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.
  */
-import { Text } from '@primer/react';
+import { useCallback, useEffect, useState } from 'react';
+import { Spinner, Text } from '@primer/react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { Box } from '@datalayer/primer-addons';
 import { datalayerTheme, DatalayerThemeProvider } from '@datalayer/core';
@@ -19,14 +20,107 @@ const queryClient = new QueryClient({
         },
     },
 });
+const AGENT_NAME = 'datalayer-assistant';
+const BASE_URL = 'http://localhost:8765';
+/**
+ * Hook to ensure the agent exists on the server
+ * Creates the agent if it doesn't exist
+ */
+function useEnsureAgent() {
+    const [agentId, setAgentId] = useState(null);
+    const [isLoading, setIsLoading] = useState(true);
+    const [error, setError] = useState(null);
+    const createAgent = useCallback(async () => {
+        try {
+            const response = await fetch(`${BASE_URL}/api/v1/agents`, {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                },
+                body: JSON.stringify({
+                    name: AGENT_NAME,
+                    description: 'Datalayer AI Assistant - Your helpful coding companion',
+                    agent_library: 'pydantic-ai',
+                    transport: 'ag-ui',
+                    model: 'openai:gpt-4o-mini',
+                    system_prompt: `You are Datalayer Assistant, a helpful AI assistant specialized in data science, Python programming, and Jupyter notebooks.
+
+You can help users with:
+- Writing and debugging Python code
+- Data analysis with pandas, numpy, and other libraries
+- Creating visualizations with matplotlib, plotly, etc.
+- Machine learning with scikit-learn, pytorch, etc.
+- General programming questions
+
+Be concise, helpful, and provide working code examples when appropriate.`,
+                }),
+            });
+            if (!response.ok) {
+                const errorData = await response
+                    .json()
+                    .catch(() => ({ detail: 'Unknown error' }));
+                throw new Error(errorData.detail || `Failed to create agent: ${response.status}`);
+            }
+            const data = await response.json();
+            console.log('[AgentRuntimeCustomExample] Agent created:', data);
+            return data.id;
+        }
+        catch (err) {
+            console.error('[AgentRuntimeCustomExample] Error creating agent:', err);
+            throw err;
+        }
+    }, []);
+    const checkOrCreateAgent = useCallback(async () => {
+        setIsLoading(true);
+        setError(null);
+        try {
+            // First, try to find existing agent by name
+            const listResponse = await fetch(`${BASE_URL}/api/v1/agents`);
+            if (listResponse.ok) {
+                const data = await listResponse.json();
+                // API may return { agents: [...] } or just [...]
+                const agents = Array.isArray(data) ? data : data.agents || [];
+                const existingAgent = agents.find((a) => a.name === AGENT_NAME);
+                if (existingAgent) {
+                    console.log('[AgentRuntimeCustomExample] Found existing agent:', existingAgent.id);
+                    setAgentId(existingAgent.id);
+                    setIsLoading(false);
+                    return;
+                }
+            }
+            // Agent doesn't exist, create it
+            console.log('[AgentRuntimeCustomExample] Creating new agent...');
+            const newAgentId = await createAgent();
+            if (newAgentId) {
+                setAgentId(newAgentId);
+            }
+            else {
+                setError('Failed to create agent');
+            }
+        }
+        catch (err) {
+            const errorMessage = err instanceof Error ? err.message : 'Failed to initialize agent';
+            setError(errorMessage);
+        }
+        finally {
+            setIsLoading(false);
+        }
+    }, [createAgent]);
+    useEffect(() => {
+        checkOrCreateAgent();
+    }, [checkOrCreateAgent]);
+    return { agentId, isLoading, error, retry: checkOrCreateAgent };
+}
 /**
  * Agent Runtime Custom Example Component
  *
- * Demonstrates the unified Chat component with Simple transport
+ * Demonstrates the unified Chat component with AG-UI transport
  * and all necessary providers:
  * - QueryClientProvider for data fetching
+ * - Auto-creates agent on the server if it doesn't exist
  */
 const AgentRuntimeCustomExample = () => {
+    const { agentId, isLoading, error, retry } = useEnsureAgent();
     return (_jsx(QueryClientProvider, { client: queryClient, children: _jsx(DatalayerThemeProvider, { theme: datalayerTheme, children: _jsxs(Box, { sx: {
                     display: 'flex',
                     flexDirection: 'column',
@@ -41,12 +135,33 @@ const AgentRuntimeCustomExample = () => {
                                     fontWeight: 'bold',
                                     display: 'block',
                                     marginBottom: 1,
-                                }, children: "Agent Runtime Custom Example" }), _jsx(Text, { sx: { fontSize: 1, color: 'fg.muted' }, children: "Interactive chat interface with AI assistance" })] }), _jsx(Box, { as: "main", sx: {
+                                }, children: "Datalayer Assistant" }), _jsx(Text, { sx: { fontSize: 1, color: 'fg.muted' }, children: "Interactive chat interface with AI assistance" })] }), _jsx(Box, { as: "main", sx: {
                             flex: 1,
                             overflow: 'hidden',
                             display: 'flex',
                             flexDirection: 'column',
-                        }, children: _jsx(Chat, { transport: "vercel-ai-jupyter", showModelSelector: true, showToolsMenu: true, height: "100%", suggestions: [
+                        }, children: isLoading ? (_jsxs(Box, { sx: {
+                                display: 'flex',
+                                flexDirection: 'column',
+                                alignItems: 'center',
+                                justifyContent: 'center',
+                                height: '100%',
+                                gap: 3,
+                            }, children: [_jsx(Spinner, { size: "large" }), _jsx(Text, { sx: { color: 'fg.muted' }, children: "Initializing agent..." })] })) : error ? (_jsxs(Box, { sx: {
+                                display: 'flex',
+                                flexDirection: 'column',
+                                alignItems: 'center',
+                                justifyContent: 'center',
+                                height: '100%',
+                                gap: 3,
+                                p: 4,
+                            }, children: [_jsx(Text, { sx: { color: 'danger.fg', fontWeight: 'bold' }, children: "Failed to initialize agent" }), _jsx(Text, { sx: { color: 'fg.muted', textAlign: 'center' }, children: error }), _jsx(Text, { as: "button", onClick: retry, sx: {
+                                        color: 'accent.fg',
+                                        cursor: 'pointer',
+                                        textDecoration: 'underline',
+                                        border: 'none',
+                                        background: 'none',
+                                    }, children: "Retry" })] })) : agentId ? (_jsx(Chat, { transport: "ag-ui", agentId: agentId, showModelSelector: true, showToolsMenu: true, height: "100%", suggestions: [
                                 {
                                     title: '👋 Say hello',
                                     message: 'Hello! What can you help me with today?',
@@ -63,6 +178,6 @@ const AgentRuntimeCustomExample = () => {
                                     title: '📊 Data analysis',
                                     message: 'How do I analyze data with pandas?',
                                 },
-                            ] }) })] }) }) }));
+                            ] })) : null })] }) }) }));
 };
 export default AgentRuntimeCustomExample;

package/lib/examples/AgentRuntimeLexicalExample.js

@@ -239,7 +239,7 @@ function LexicalWithChat({ content, serviceManager, }) {
                     color: 'danger.fg',
                     borderRadius: 2,
                     maxWidth: 300,
-                }, children: [_jsx("strong", { children: "Error:" }), " ", error] })), isReady && (_jsx(ChatFloating, { endpoint: AG_UI_ENDPOINT, title: "Lexical AI Agent Runtime", description: "Hi! I can help you edit documents. Try: 'Insert a heading', 'Add a code block', or 'Create a list'", defaultOpen: true, defaultViewMode: "panel", position: "bottom-right", brandColor: "#7c3aed", tools: tools, useStore: false, suggestions: [
+                }, children: [_jsx("strong", { children: "Error:" }), " ", error] })), isReady && (_jsx(ChatFloating, { endpoint: AG_UI_ENDPOINT, title: "Lexical AI Agent Runtime", description: "Hi! I can help you edit documents. Try: 'Insert a heading', 'Add a code block', or 'Create a list'", defaultOpen: true, defaultViewMode: "panel", position: "bottom-right", brandColor: "#7c3aed", tools: tools, useStore: false, showModelSelector: true, showToolsMenu: true, suggestions: [
                     {
                         title: 'Insert heading',
                         message: 'Insert a heading that says "Welcome"',

package/lib/examples/components/AgentConfiguration.d.ts

@@ -1,6 +1,28 @@
 import React from 'react';
 import type { Agent } from '../stores/examplesStore';
 import type { Transport, Extension } from '../../components/chat';
+/**
+ * MCP Server Tool type
+ */
+export interface MCPServerTool {
+    name: string;
+    description?: string;
+    enabled: boolean;
+}
+/**
+ * MCP Server configuration from backend
+ */
+export interface MCPServerConfig {
+    id: string;
+    name: string;
+    url?: string;
+    enabled: boolean;
+    tools: MCPServerTool[];
+    command?: string;
+    args?: string[];
+    isAvailable?: boolean;
+    transport?: string;
+}
 type AgentLibrary = 'pydantic-ai' | 'langchain' | 'jupyter-ai';
 export type { AgentLibrary };
 export type { Transport };

package/lib/examples/components/AgentConfiguration.js

@@ -1,5 +1,12 @@
 import { jsx as _jsx, jsxs as _jsxs } from "react/jsx-runtime";
-import { Text, TextInput, Button, FormControl, Select, CheckboxGroup, Checkbox, Spinner, Flash, } from '@primer/react';
+/*
+ * Copyright (c) 2025-2026 Datalayer, Inc.
+ * Distributed under the terms of the Modified BSD License.
+ */
+import React from 'react';
+import { Text, TextInput, Button, FormControl, Select, Checkbox, Spinner, Flash, Label, ActionList, } from '@primer/react';
+import { CheckIcon, XIcon, ToolsIcon } from '@primer/octicons-react';
+import { useQuery } from '@tanstack/react-query';
 import { Box } from '@datalayer/primer-addons';
 const AGENT_LIBRARIES = [
     {
@@ -65,6 +72,21 @@ const EXTENSIONS = [
  * Form for configuring agent connection settings.
  */
 export const AgentConfiguration = ({ agentLibrary, transport, extensions, wsUrl, baseUrl, agentName, agents, selectedAgentId, isCreatingAgent = false, createError = null, onAgentLibraryChange, onTransportChange, onExtensionsChange, onWsUrlChange, onBaseUrlChange, onAgentNameChange, onAgentSelect, onConnect, }) => {
+    // Fetch MCP servers configuration from the backend
+    const configQuery = useQuery({
+        queryKey: ['agent-config', baseUrl],
+        queryFn: async () => {
+            const response = await fetch(`${baseUrl}/api/v1/configure`);
+            if (!response.ok) {
+                throw new Error('Failed to fetch configuration');
+            }
+            return response.json();
+        },
+        enabled: !!baseUrl,
+        staleTime: 1000 * 60 * 5, // 5 minutes
+        retry: 1,
+    });
+    const mcpServers = configQuery.data?.mcpServers || [];
     // Determine which extensions are enabled based on transport
     const isExtensionEnabled = (ext) => {
         if (selectedAgentId !== 'new-agent')
@@ -95,7 +117,7 @@ export const AgentConfiguration = ({ agentLibrary, transport, extensions, wsUrl,
                     fontWeight: 'bold',
                     display: 'block',
                     marginBottom: 3,
-                }, children: "Connection Settings" }), _jsxs(FormControl, { sx: { marginBottom: 3 }, children: [_jsx(FormControl.Label, { children: "Available Agents" }), _jsxs(Select, { value: selectedAgentId, onChange: e => onAgentSelect(e.target.value), sx: { width: '100%' }, children: [_jsx(Select.Option, { value: "new-agent", children: "+ New Agent..." }), agents.map(agent => (_jsxs(Select.Option, { value: agent.id, children: [agent.status === 'running' && '● ', agent.name] }, agent.id)))] }), _jsx(FormControl.Caption, { children: selectedAgentId === 'new-agent'
+                }, children: "Create a new Agent" }), _jsxs(FormControl, { sx: { marginBottom: 3 }, children: [_jsx(FormControl.Label, { children: "Available Agents" }), _jsxs(Select, { value: selectedAgentId, onChange: e => onAgentSelect(e.target.value), sx: { width: '100%' }, children: [_jsx(Select.Option, { value: "new-agent", children: "+ New Agent..." }), agents.map(agent => (_jsxs(Select.Option, { value: agent.id, children: [agent.status === 'running' && '● ', agent.name] }, agent.id)))] }), _jsx(FormControl.Caption, { children: selectedAgentId === 'new-agent'
                             ? 'Configure a new custom agent'
                             : 'Selected agent - form fields below are disabled' })] }), _jsxs(Box, { sx: { display: 'flex', gap: 3, marginBottom: 3 }, children: [_jsxs(FormControl, { sx: { flex: 1 }, children: [_jsx(FormControl.Label, { children: "Agent Library" }), _jsx(Select, { value: agentLibrary, onChange: e => onAgentLibraryChange(e.target.value), disabled: selectedAgentId !== 'new-agent', sx: { width: '100%' }, children: AGENT_LIBRARIES.map(lib => (_jsxs(Select.Option, { value: lib.value, disabled: lib.disabled, children: [lib.label, lib.disabled && ' (Coming Soon)'] }, lib.value))) })] }), _jsxs(FormControl, { sx: { flex: 1 }, children: [_jsx(FormControl.Label, { children: "Transport" }), _jsx(Select, { value: transport, onChange: e => onTransportChange(e.target.value), disabled: selectedAgentId !== 'new-agent', sx: { width: '100%' }, children: TRANSPORTS.map(t => (_jsx(Select.Option, { value: t.value, children: t.label }, t.value))) })] }), _jsxs(FormControl, { sx: { flex: 1 }, children: [_jsx(FormControl.Label, { children: "Extensions" }), _jsx(Box, { sx: { display: 'flex', flexDirection: 'column', gap: 2 }, children: EXTENSIONS.map(ext => (_jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 2 }, children: [_jsx(Checkbox, { value: ext.value, checked: extensions.includes(ext.value), disabled: !isExtensionEnabled(ext.value), onChange: e => handleExtensionChange(ext.value, e.target.checked) }), _jsx(Text, { children: ext.label })] }, ext.value))) })] })] }), _jsxs(FormControl, { sx: { marginBottom: 3 }, children: [_jsx(FormControl.Label, { children: transport === 'acp' ? 'WebSocket URL' : 'Base URL' }), _jsx(TextInput, { value: transport === 'acp' ? wsUrl : baseUrl, onChange: e => transport === 'acp'
                             ? onWsUrlChange(e.target.value)
@@ -103,7 +125,21 @@ export const AgentConfiguration = ({ agentLibrary, transport, extensions, wsUrl,
                             ? 'ws://localhost:8000/api/v1/acp/ws'
                             : 'http://localhost:8000', sx: { width: '100%' } }), _jsx(FormControl.Caption, { children: transport === 'acp'
                             ? 'The WebSocket endpoint of your agent-runtimes server'
-                            : 'The base URL of your agent-runtimes server' })] }), _jsxs(FormControl, { sx: { marginBottom: 3 }, children: [_jsx(FormControl.Label, { children: "Agent Name" }), _jsx(TextInput, { value: agentName, onChange: e => onAgentNameChange(e.target.value), disabled: selectedAgentId !== 'new-agent', placeholder: "demo-agent", sx: { width: '100%' } }), _jsx(FormControl.Caption, { children: "The name of the agent to connect to" })] }), _jsxs(CheckboxGroup, { sx: { marginBottom: 3 }, disabled: true, children: [_jsx(CheckboxGroup.Label, { children: "MCP Servers (Coming Soon)" }), _jsx(CheckboxGroup.Caption, { children: "Select MCP servers to connect to" }), _jsxs(FormControl, { disabled: true, children: [_jsx(Checkbox, { value: "github", defaultChecked: true, disabled: true }), _jsx(FormControl.Label, { children: "GitHub" })] }), _jsxs(FormControl, { disabled: true, children: [_jsx(Checkbox, { value: "anaconda", disabled: true }), _jsx(FormControl.Label, { children: "Anaconda" })] }), _jsxs(FormControl, { disabled: true, children: [_jsx(Checkbox, { value: "tavily", disabled: true }), _jsx(FormControl.Label, { children: "Tavily" })] })] }), createError && (_jsx(Flash, { variant: "danger", sx: { marginBottom: 3 }, children: createError })), _jsx(Button, { variant: "primary", onClick: onConnect, disabled: isCreatingAgent ||
+                            : 'The base URL of your agent-runtimes server' })] }), _jsxs(FormControl, { sx: { marginBottom: 3 }, children: [_jsx(FormControl.Label, { children: "Agent Name" }), _jsx(TextInput, { value: agentName, onChange: e => onAgentNameChange(e.target.value), disabled: selectedAgentId !== 'new-agent', placeholder: "demo-agent", sx: { width: '100%' } }), _jsx(FormControl.Caption, { children: "The name of the agent to connect to" })] }), _jsxs(Box, { sx: {
+                    marginBottom: 3,
+                    padding: 3,
+                    border: '1px solid',
+                    borderColor: 'border.default',
+                    borderRadius: 2,
+                    backgroundColor: 'canvas.default',
+                }, children: [_jsxs(Box, { sx: {
+                            display: 'flex',
+                            alignItems: 'center',
+                            gap: 2,
+                            marginBottom: 2,
+                        }, children: [_jsx(ToolsIcon, { size: 16 }), _jsx(Text, { sx: { fontSize: 1, fontWeight: 'bold' }, children: "MCP Servers" }), configQuery.isLoading && _jsx(Spinner, { size: "small" })] }), configQuery.isError && (_jsx(Flash, { variant: "warning", sx: { marginBottom: 2 }, children: _jsx(Text, { sx: { fontSize: 0 }, children: "Unable to fetch MCP servers. Check that the server is running." }) })), mcpServers.length === 0 &&
+                        !configQuery.isLoading &&
+                        !configQuery.isError && (_jsx(Text, { sx: { fontSize: 0, color: 'fg.muted' }, children: "No MCP servers configured." })), mcpServers.length > 0 && (_jsx(ActionList, { children: mcpServers.map((server, index) => (_jsxs(React.Fragment, { children: [index > 0 && _jsx(ActionList.Divider, {}), _jsxs(ActionList.Item, { disabled: true, children: [_jsx(ActionList.LeadingVisual, { children: server.isAvailable ? (_jsx(CheckIcon, { size: 16 })) : (_jsx(XIcon, { size: 16 })) }), _jsxs(Box, { sx: { display: 'flex', flexDirection: 'column', gap: 1 }, children: [_jsxs(Box, { sx: { display: 'flex', alignItems: 'center', gap: 2 }, children: [_jsx(Text, { sx: { fontWeight: 'semibold' }, children: server.name }), _jsx(Label, { variant: server.isAvailable ? 'success' : 'secondary', size: "small", children: server.isAvailable ? 'Available' : 'Not Available' })] }), server.tools.length > 0 && (_jsxs(Text, { sx: { fontSize: 0, color: 'fg.muted' }, children: ["Tools: ", server.tools.map(t => t.name).join(', ')] }))] })] })] }, server.id))) }))] }), createError && (_jsx(Flash, { variant: "danger", sx: { marginBottom: 3 }, children: createError })), _jsx(Button, { variant: "primary", onClick: onConnect, disabled: isCreatingAgent ||
                     !agentName ||
                     (transport === 'acp' ? !wsUrl : !baseUrl), sx: { width: '100%' }, children: isCreatingAgent ? (_jsxs(Box, { sx: {
                         display: 'flex',

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

@@ -8,4 +8,4 @@ export { TimeTravel } from './TimeTravel';
 export { LexicalEditor } from './LexicalEditor';
 export { Rating } from './Rating';
 export { AgentConfiguration, AGENT_LIBRARIES, TRANSPORTS, EXTENSIONS, } from './AgentConfiguration';
-export type { AgentLibrary, Transport, Extension } from './AgentConfiguration';
+export type { AgentLibrary, Transport, Extension, MCPServerConfig, MCPServerTool, } from './AgentConfiguration';

package/lib/hooks/useNotebookAIAgent.d.ts

@@ -1,8 +1,8 @@
-import type { IAIAgent } from '../models';
+import type { Agent } from '../state';
 /**
  * Get the document AI Agent if any.
  *
  * It handles checking the AI Agent is alive so it should only be use once per document.
  */
-export declare function useNotebookAIAgent(notebookId: string): IAIAgent | undefined;
+export declare function useNotebookAIAgent(notebookId: string): Agent | undefined;
 export default useNotebookAIAgent;

package/lib/hooks/useNotebookAIAgent.js

@@ -3,7 +3,7 @@
  * Distributed under the terms of the Modified BSD License.
  */
 import { useEffect } from 'react';
-import { useAIAgentStore } from '../state';
+import { useAgentStore } from '../state';
 import { useAIAgents } from './useAgents';
 /**
  * Get the document AI Agent if any.
@@ -12,7 +12,10 @@ import { useAIAgents } from './useAgents';
  */
 export function useNotebookAIAgent(notebookId) {
     const { getAIAgent } = useAIAgents();
-    const { aiAgents, addAIAgent, deleteAIAgent, updateAIAgent } = useAIAgentStore();
+    const agents = useAgentStore(state => state.agents);
+    const upsertAgent = useAgentStore(state => state.upsertAgent);
+    const deleteAgent = useAgentStore(state => state.deleteAgent);
+    const getAgentById = useAgentStore(state => state.getAgentById);
     // Check AI Agent is alive.
     useEffect(() => {
         let abortController;
@@ -23,25 +26,39 @@ export function useNotebookAIAgent(notebookId) {
                     signal: abortController.signal,
                 });
                 if (!response.success) {
-                    deleteAIAgent(notebookId);
+                    deleteAgent(notebookId);
                     return;
                 }
-                const currentAgent = aiAgents.find(agent => agent.documentId === notebookId);
+                const currentAgent = getAgentById(notebookId);
                 const runtimeId = response.agent.runtime?.id;
                 if (currentAgent) {
+                    // Update existing agent if runtime changed
                     if (currentAgent.runtimeId !== runtimeId) {
-                        updateAIAgent(notebookId, runtimeId);
+                        upsertAgent({
+                            id: notebookId,
+                            baseUrl: currentAgent.baseUrl,
+                            transport: currentAgent.transport,
+                            runtimeId,
+                            status: 'running',
+                        });
                     }
                 }
                 else {
-                    addAIAgent({
+                    // Add new agent - assume it's a notebook agent with document tracking
+                    upsertAgent({
+                        id: notebookId,
+                        name: `Notebook ${notebookId}`,
+                        description: 'AI agent for notebook',
+                        baseUrl: '', // Will be set by the actual agent implementation
+                        transport: 'vercel-ai', // Default transport for notebook agents
                         documentId: notebookId,
                         runtimeId,
+                        status: 'running',
                     });
                 }
             }
             catch (r) {
-                deleteAIAgent(notebookId);
+                deleteAgent(notebookId);
             }
         };
         const refreshInterval = setInterval(refreshAIAgent, 60_000);
@@ -49,8 +66,8 @@ export function useNotebookAIAgent(notebookId) {
             abortController?.abort('Component unmounted');
             clearInterval(refreshInterval);
         };
-    }, [aiAgents, notebookId]);
-    const aiAgent = aiAgents.find(aiAgent => aiAgent.documentId === notebookId);
+    }, [agents, notebookId, getAIAgent, deleteAgent, getAgentById, upsertAgent]);
+    const aiAgent = getAgentById(notebookId);
     return aiAgent;
 }
 export default useNotebookAIAgent;

package/lib/index.d.ts

@@ -1 +1,2 @@
 export * from './components';
+export * from './state';

package/lib/index.js

@@ -3,3 +3,4 @@
  * Distributed under the terms of the Modified BSD License.
  */
 export * from './components';
+export * from './state';

package/lib/state/substates/AIAgentState.d.ts

@@ -1,11 +1,81 @@
-import type { IAIAgent } from '../../models';
-export type AIAgentState = {
-    aiAgents: readonly IAIAgent[];
-    addAIAgent: (aiAgent: IAIAgent) => void;
-    deleteAIAgent: (documentId: string) => void;
-    updateAIAgent: (documentId: string, runtimeId?: string) => void;
+import type { Transport } from '../../components/chat/components/Chat';
+export type AgentStatus = 'running' | 'paused' | 'initializing' | 'error';
+export type { Transport };
+/**
+ * Unified Agent model combining runtime tracking and UI state
+ */
+export interface Agent {
+    /** Unique agent identifier */
+    id: string;
+    /** Display name */
+    name: string;
+    /** Agent description */
+    description: string;
+    /** Base URL for the agent (for Jupyter: baseUrl, for FastAPI: baseUrl) */
+    baseUrl: string;
+    /** Transport protocol used */
+    transport: Transport;
+    /** Current status */
+    status: AgentStatus;
+    /** Last error message if status is 'error' */
+    error?: string | null;
+    /** Timestamp of last update */
+    lastUpdated: number;
+    /** Document ID (for document-based agents) */
+    documentId?: string;
+    /** Runtime ID (for Jupyter kernel-based agents) */
+    runtimeId?: string;
+    /** Author name (optional, for display) */
+    author?: string;
+    /** Avatar URL (optional, for display) */
+    avatarUrl?: string;
+}
+export type AgentState = {
+    /** All registered agents */
+    agents: readonly Agent[];
+    /** Add or update an agent */
+    upsertAgent: (agent: Partial<Agent> & {
+        id: string;
+        baseUrl: string;
+        transport: Transport;
+    }) => void;
+    /** Get agent by ID */
+    getAgentById: (id: string) => Agent | undefined;
+    /** Get agent by baseUrl and transport */
+    getAgentByUrl: (baseUrl: string, transport: Transport) => Agent | undefined;
+    /** Update agent status */
+    updateAgentStatus: (id: string, status: AgentStatus, error?: string | null) => void;
+    /** Toggle agent status between running/paused */
+    toggleAgentStatus: (id: string) => void;
+    /** Delete an agent */
+    deleteAgent: (id: string) => void;
+    /** Clear all agents */
+    clearAgents: () => void;
 };
-export declare const aiAgentStore: import("zustand").StoreApi<AIAgentState>;
-export declare function useAIAgentStore(): AIAgentState;
-export declare function useAIAgentStore<T>(selector: (state: AIAgentState) => T): T;
-export default useAIAgentStore;
+export declare const agentStore: Omit<import("zustand").StoreApi<AgentState>, "persist"> & {
+    persist: {
+        setOptions: (options: Partial<import("zustand/middleware").PersistOptions<AgentState, unknown>>) => void;
+        clearStorage: () => void;
+        rehydrate: () => Promise<void> | void;
+        hasHydrated: () => boolean;
+        onHydrate: (fn: (state: AgentState) => void) => () => void;
+        onFinishHydration: (fn: (state: AgentState) => void) => () => void;
+        getOptions: () => Partial<import("zustand/middleware").PersistOptions<AgentState, unknown>>;
+    };
+};
+export declare function useAgentStore(): AgentState;
+export declare function useAgentStore<T>(selector: (state: AgentState) => T): T;
+export type AIAgentState = AgentState;
+export declare const aiAgentStore: Omit<import("zustand").StoreApi<AgentState>, "persist"> & {
+    persist: {
+        setOptions: (options: Partial<import("zustand/middleware").PersistOptions<AgentState, unknown>>) => void;
+        clearStorage: () => void;
+        rehydrate: () => Promise<void> | void;
+        hasHydrated: () => boolean;
+        onHydrate: (fn: (state: AgentState) => void) => () => void;
+        onFinishHydration: (fn: (state: AgentState) => void) => () => void;
+        getOptions: () => Partial<import("zustand/middleware").PersistOptions<AgentState, unknown>>;
+    };
+};
+export declare const useAIAgentStore: typeof useAgentStore;
+export default useAgentStore;