@jupyterlite/ai 0.12.0 → 0.14.0

package/lib/agent.d.ts

@@ -1,6 +1,7 @@
 import { ISignal } from '@lumino/signaling';
 import { type Tool } from 'ai';
 import { ISecretsManager } from 'jupyter-secrets-manager';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import { AISettingsModel } from './models/settings-model';
 import type { IProviderRegistry } from './tokens';
 import { ISkillRegistry, ITool, IToolRegistry, ITokenUsage } from './tokens';
@@ -22,7 +23,7 @@ export declare namespace AgentManagerFactory {
         /**
          * The token used to request the secrets manager.
          */
-        token: symbol;
+        token: symbol | null;
     }
 }
 export declare class AgentManagerFactory {
@@ -92,7 +93,7 @@ export interface IAgentEventTypeMap {
     tool_call_complete: {
         callId: string;
         toolName: string;
-        output: string;
+        outputData: unknown;
         isError: boolean;
     };
     tool_approval_request: {
@@ -148,6 +149,10 @@ export interface IAgentManagerOptions {
      * Initial token usage.
      */
     tokenUsage?: ITokenUsage;
+    /**
+     * JupyterLab render mime registry for discovering supported MIME types.
+     */
+    renderMimeRegistry?: IRenderMimeRegistry;
 }
 /**
  * Manages the AI agent lifecycle and execution loop.
@@ -245,6 +250,10 @@ export declare class AgentManager {
      * Prepare model, tools, and settings needed to (re)build the agent.
      */
     private _prepareAgentConfig;
+    /**
+     * Build the runtime tool map used by the agent.
+     */
+    private _buildRuntimeTools;
     /**
      * Rebuild the agent using cached resources and the current skills snapshot.
      */
@@ -264,6 +273,14 @@ export declare class AgentManager {
      * Handles tool-result stream parts.
      */
     private _handleToolResult;
+    /**
+     * Handles tool-error stream parts.
+     */
+    private _handleToolError;
+    /**
+     * Handles tool-output-denied stream parts.
+     */
+    private _handleToolOutputDenied;
     /**
      * Handles tool-approval-request stream parts.
      */
@@ -296,6 +313,22 @@ export declare class AgentManager {
      * @returns The enhanced system prompt with dynamic additions
      */
     private _getEnhancedSystemPrompt;
+    /**
+     * Build an instruction line describing MIME types supported by this session.
+     */
+    private _getSupportedMimeTypesInstruction;
+    /**
+     * Sanitizes history to ensure it's in a valid state in case of abort or error.
+     */
+    private _sanitizeHistory;
+    /**
+     * Extracts tool call IDs from a message
+     */
+    private _getToolCallIds;
+    /**
+     * Checks if a tool message contains results for all specified tool call IDs
+     */
+    private _matchesAllToolCalls;
     private _settingsModel;
     private _toolRegistry?;
     private _providerRegistry?;
@@ -312,6 +345,7 @@ export declare class AgentManager {
     private _activeProvider;
     private _activeProviderChanged;
     private _skills;
+    private _renderMimeRegistry?;
     private _initQueue;
     private _agentConfig;
     private _pendingApprovals;

package/lib/agent.js

@@ -2,6 +2,7 @@ import { Signal } from '@lumino/signaling';
 import { ToolLoopAgent, stepCountIs } from 'ai';
 import { createMCPClient } from '@ai-sdk/mcp';
 import { createModel } from './providers/models';
+import { createProviderTools } from './providers/provider-tools';
 import { SECRETS_NAMESPACE } from './tokens';
 export class AgentManagerFactory {
     constructor(options) {
@@ -20,6 +21,10 @@ export class AgentManagerFactory {
         this._initializeAgents().catch(error => console.warn('Failed to initialize agent in constructor:', error));
         // Listen for settings changes
         this._settingsModel.stateChanged.connect(this._onSettingsChanged, this);
+        // Disable the secrets manager if the token is empty.
+        if (!options.token) {
+            this._secretsManager = undefined;
+        }
     }
     createAgent(options) {
         const agentManager = new AgentManager({
@@ -28,6 +33,16 @@ export class AgentManagerFactory {
             secretsManager: this._secretsManager
         });
         this._agentManagers.push(agentManager);
+        // New chats can be created before MCP setup finishes.
+        // Reinitialize them with connected MCP tools once it does.
+        this._initQueue
+            .then(() => this.getMCPTools())
+            .then(mcpTools => {
+            if (Object.keys(mcpTools).length > 0) {
+                agentManager.initializeAgent(mcpTools);
+            }
+        })
+            .catch(error => console.warn('Failed to pass MCP tools to new agent:', error));
         return agentManager;
     }
     /**
@@ -180,6 +195,7 @@ export class AgentManager {
         this._tokenUsageChanged = new Signal(this);
         this._skills = [];
         this._agentConfig = null;
+        this._renderMimeRegistry = options.renderMimeRegistry;
         this.activeProvider =
             options.activeProvider ?? this._settingsModel.config.defaultProvider;
         // Initialize selected tools to all available tools by default
@@ -372,7 +388,7 @@ export class AgentManager {
                 this._updateTokenUsage(await result.usage);
                 // Add response messages to history
                 if (responseMessages.messages?.length) {
-                    this._history.push(...responseMessages.messages);
+                    this._history.push(...Private.sanitizeModelMessages(responseMessages.messages));
                 }
                 // Add approval response if processed
                 if (streamResult.approvalResponse) {
@@ -400,6 +416,9 @@ export class AgentManager {
                     data: { error: error }
                 });
             }
+            // After an error (including AbortError), sanitize the history
+            // to remove any trailing assistant messages without tool results
+            this._sanitizeHistory();
         }
         finally {
             this._controller = null;
@@ -454,18 +473,27 @@ export class AgentManager {
             this._mcpTools = mcpTools;
         }
         const model = await this._createModel();
-        const shouldUseTools = !!(config.toolsEnabled &&
-            this._selectedToolNames.length > 0 &&
-            this._toolRegistry &&
-            Object.keys(this._toolRegistry.tools).length > 0 &&
-            this._supportsToolCalling());
-        const tools = shouldUseTools
-            ? { ...this.selectedAgentTools, ...this._mcpTools }
-            : this._mcpTools;
+        const supportsToolCalling = this._supportsToolCalling();
+        const canUseTools = config.toolsEnabled && supportsToolCalling;
+        const hasFunctionToolRegistry = !!(this._toolRegistry && Object.keys(this._toolRegistry.tools).length > 0);
+        const selectedFunctionTools = canUseTools && hasFunctionToolRegistry ? this.selectedAgentTools : {};
+        const functionTools = canUseTools
+            ? { ...selectedFunctionTools, ...this._mcpTools }
+            : {};
         const activeProviderConfig = this._settingsModel.getProvider(this._activeProvider);
+        const activeProviderInfo = activeProviderConfig && this._providerRegistry
+            ? this._providerRegistry.getProviderInfo(activeProviderConfig.provider)
+            : null;
         const temperature = activeProviderConfig?.parameters?.temperature ?? DEFAULT_TEMPERATURE;
         const maxTokens = activeProviderConfig?.parameters?.maxOutputTokens;
         const maxTurns = activeProviderConfig?.parameters?.maxTurns ?? DEFAULT_MAX_TURNS;
+        const tools = this._buildRuntimeTools({
+            providerInfo: activeProviderInfo,
+            customSettings: activeProviderConfig?.customSettings,
+            functionTools,
+            includeProviderTools: canUseTools
+        });
+        const shouldUseTools = canUseTools && Object.keys(tools).length > 0;
         this._agentConfig = {
             model,
             tools,
@@ -476,6 +504,22 @@ export class AgentManager {
             shouldUseTools
         };
     }
+    /**
+     * Build the runtime tool map used by the agent.
+     */
+    _buildRuntimeTools(options) {
+        const providerTools = options.includeProviderTools
+            ? createProviderTools({
+                providerInfo: options.providerInfo,
+                customSettings: options.customSettings,
+                hasFunctionTools: Object.keys(options.functionTools).length > 0
+            })
+            : {};
+        return {
+            ...providerTools,
+            ...options.functionTools
+        };
+    }
     /**
      * Rebuild the agent using cached resources and the current skills snapshot.
      */
@@ -485,9 +529,21 @@ export class AgentManager {
             return;
         }
         const { model, tools, temperature, maxOutputTokens, maxTurns, baseSystemPrompt, shouldUseTools } = this._agentConfig;
-        const instructions = shouldUseTools
-            ? this._getEnhancedSystemPrompt(baseSystemPrompt)
+        const baseInstructions = shouldUseTools
+            ? this._getEnhancedSystemPrompt(baseSystemPrompt, tools)
             : baseSystemPrompt || 'You are a helpful assistant.';
+        const richOutputWorkflowInstruction = shouldUseTools
+            ? '- When the user asks for visual or rich outputs, prefer running code/commands that produce those outputs and describe that they will be rendered in chat.'
+            : '- When tools are unavailable, explain the limitation clearly and provide concrete steps the user can run to produce the desired rich outputs.';
+        const supportedMimeTypesInstruction = this._getSupportedMimeTypesInstruction();
+        const instructions = `${baseInstructions}
+
+RICH OUTPUT RENDERING:
+- The chat UI can render rich MIME outputs as separate assistant messages.
+- ${supportedMimeTypesInstruction}
+- Use only MIME types from the supported list when creating MIME bundles. Do not invent MIME keys.
+- Do not claim that you cannot display maps, images, or rich outputs in chat.
+${richOutputWorkflowInstruction}`;
         this._agent = new ToolLoopAgent({
             model,
             instructions,
@@ -546,6 +602,12 @@ export class AgentManager {
                 case 'tool-result':
                     this._handleToolResult(part);
                     break;
+                case 'tool-error':
+                    this._handleToolError(part);
+                    break;
+                case 'tool-output-denied':
+                    this._handleToolOutputDenied(part);
+                    break;
                 case 'tool-approval-request':
                     // Complete current message before approval
                     if (currentMessageId && fullResponse) {
@@ -579,9 +641,6 @@ export class AgentManager {
      * Handles tool-result stream parts.
      */
     _handleToolResult(part) {
-        const output = typeof part.output === 'string'
-            ? part.output
-            : JSON.stringify(part.output, null, 2);
         const isError = typeof part.output === 'object' &&
             part.output !== null &&
             'success' in part.output &&
@@ -591,11 +650,44 @@ export class AgentManager {
             data: {
                 callId: part.toolCallId,
                 toolName: part.toolName,
-                output,
+                outputData: part.output,
                 isError
             }
         });
     }
+    /**
+     * Handles tool-error stream parts.
+     */
+    _handleToolError(part) {
+        const output = typeof part.error === 'string'
+            ? part.error
+            : part.error instanceof Error
+                ? part.error.message
+                : JSON.stringify(part.error, null, 2);
+        this._agentEvent.emit({
+            type: 'tool_call_complete',
+            data: {
+                callId: part.toolCallId,
+                toolName: part.toolName,
+                outputData: output,
+                isError: true
+            }
+        });
+    }
+    /**
+     * Handles tool-output-denied stream parts.
+     */
+    _handleToolOutputDenied(part) {
+        this._agentEvent.emit({
+            type: 'tool_call_complete',
+            data: {
+                callId: part.toolCallId,
+                toolName: part.toolName,
+                outputData: 'Tool output was denied.',
+                isError: true
+            }
+        });
+    }
     /**
      * Handles tool-approval-request stream parts.
      */
@@ -681,8 +773,16 @@ export class AgentManager {
         const baseURL = activeProviderConfig.baseURL;
         let apiKey;
         if (this._secretsManager && this._settingsModel.config.useSecretsManager) {
-            apiKey =
-                (await this._secretsManager.get(Private.getToken(), SECRETS_NAMESPACE, `${provider}:apiKey`))?.value ?? '';
+            const token = Private.getToken();
+            if (!token) {
+                // This should never happen, the secrets manager should be disabled.
+                console.error('@jupyterlite/ai::AgentManager error: the settings manager token is not set.\nYou should disable the the secrets manager from the AI settings.');
+                apiKey = '';
+            }
+            else {
+                apiKey =
+                    (await this._secretsManager.get(token, SECRETS_NAMESPACE, `${provider}:apiKey`))?.value ?? '';
+            }
         }
         else {
             apiKey = this._settingsModel.getApiKey(activeProviderConfig.id);
@@ -699,12 +799,11 @@ export class AgentManager {
      * @param baseSystemPrompt The base system prompt from settings
      * @returns The enhanced system prompt with dynamic additions
      */
-    _getEnhancedSystemPrompt(baseSystemPrompt) {
-        if (this._skills.length === 0) {
-            return baseSystemPrompt;
-        }
-        const lines = this._skills.map(skill => `- ${skill.name}: ${skill.description}`);
-        const skillsPrompt = `
+    _getEnhancedSystemPrompt(baseSystemPrompt, tools) {
+        let prompt = baseSystemPrompt;
+        if (this._skills.length > 0) {
+            const lines = this._skills.map(skill => `- ${skill.name}: ${skill.description}`);
+            const skillsPrompt = `
 
 AGENT SKILLS:
 Skills are provided via the skills registry and accessed through tools (not commands).
@@ -716,7 +815,116 @@ If the load_skill result includes a non-empty "resources" array, those are bundl
 AVAILABLE SKILLS (preloaded snapshot):
 ${lines.join('\n')}
 `;
-        return baseSystemPrompt + skillsPrompt;
+            prompt += skillsPrompt;
+        }
+        const toolNames = new Set(Object.keys(tools));
+        const hasBrowserFetch = toolNames.has('browser_fetch');
+        const hasWebFetch = toolNames.has('web_fetch');
+        const hasWebSearch = toolNames.has('web_search');
+        if (hasBrowserFetch || hasWebFetch || hasWebSearch) {
+            const webRetrievalPrompt = `
+
+WEB RETRIEVAL POLICY:
+- If the user asks about a specific URL and browser_fetch is available, call browser_fetch first for that URL.
+- If browser_fetch fails due to CORS/network/access, try web_fetch (if available) for that same URL.
+- If web_fetch fails with access/policy errors (for example: url_not_accessible or url_not_allowed) and browser_fetch is available, you MUST call browser_fetch for that same URL before searching.
+- If either fetch method fails with temporary access/network issues (for example: network_or_cors), try the other fetch method if available before searching.
+- Only fall back to web_search after both fetch methods fail or are unavailable.
+- If the user explicitly asks to inspect one exact URL, do not skip directly to search unless both fetch methods fail or are unavailable.
+- In your final response, state which retrieval method succeeded (browser_fetch, web_fetch, or web_search) and mention relevant limitations.
+`;
+            prompt += webRetrievalPrompt;
+        }
+        return prompt;
+    }
+    /**
+     * Build an instruction line describing MIME types supported by this session.
+     */
+    _getSupportedMimeTypesInstruction() {
+        const mimeTypes = this._renderMimeRegistry?.mimeTypes ?? [];
+        const safeMimeTypes = mimeTypes.filter(mimeType => {
+            const factory = this._renderMimeRegistry?.getFactory(mimeType);
+            return !!factory?.safe;
+        });
+        if (safeMimeTypes.length === 0) {
+            return 'Supported MIME types are determined by the active JupyterLab renderers in this session.';
+        }
+        return `Supported MIME types in this session: ${safeMimeTypes.join(', ')}`;
+    }
+    /**
+     * Sanitizes history to ensure it's in a valid state in case of abort or error.
+     */
+    _sanitizeHistory() {
+        if (this._history.length === 0) {
+            return;
+        }
+        const newHistory = [];
+        for (let i = 0; i < this._history.length; i++) {
+            const msg = this._history[i];
+            if (msg.role === 'assistant') {
+                const toolCallIds = this._getToolCallIds(msg);
+                if (toolCallIds.length > 0) {
+                    // Find if there's a following tool message with results for these calls
+                    const nextMsg = this._history[i + 1];
+                    if (nextMsg &&
+                        nextMsg.role === 'tool' &&
+                        this._matchesAllToolCalls(nextMsg, toolCallIds)) {
+                        newHistory.push(msg);
+                    }
+                    else {
+                        // Message has unmatched tool calls drop it and everything after it
+                        break;
+                    }
+                }
+                else {
+                    newHistory.push(msg);
+                }
+            }
+            else if (msg.role === 'tool') {
+                // Tool messages are valid if they were preceded by a valid assistant message
+                newHistory.push(msg);
+            }
+            else {
+                newHistory.push(msg);
+            }
+        }
+        this._history = newHistory;
+    }
+    /**
+     * Extracts tool call IDs from a message
+     */
+    _getToolCallIds(message) {
+        const ids = [];
+        // Check content array for tool-call parts
+        if (Array.isArray(message.content)) {
+            for (const part of message.content) {
+                if (typeof part === 'object' &&
+                    part !== null &&
+                    'type' in part &&
+                    part.type === 'tool-call') {
+                    ids.push(part.toolCallId);
+                }
+            }
+        }
+        return ids;
+    }
+    /**
+     * Checks if a tool message contains results for all specified tool call IDs
+     */
+    _matchesAllToolCalls(message, callIds) {
+        if (message.role !== 'tool' || !Array.isArray(message.content)) {
+            return false;
+        }
+        const resultIds = new Set();
+        for (const part of message.content) {
+            if (typeof part === 'object' &&
+                part !== null &&
+                'type' in part &&
+                part.type === 'tool-result') {
+                resultIds.add(part.toolCallId);
+            }
+        }
+        return callIds.every(id => resultIds.has(id));
     }
     // Private attributes
     _settingsModel;
@@ -735,12 +943,29 @@ ${lines.join('\n')}
     _activeProvider = '';
     _activeProviderChanged = new Signal(this);
     _skills;
+    _renderMimeRegistry;
     _initQueue = Promise.resolve();
     _agentConfig;
     _pendingApprovals = new Map();
 }
 var Private;
 (function (Private) {
+    /**
+     * Keep only serializable messages by doing a JSON round-trip.
+     * Messages that cannot be serialized are dropped.
+     */
+    Private.sanitizeModelMessages = (messages) => {
+        const sanitized = [];
+        for (const message of messages) {
+            try {
+                sanitized.push(JSON.parse(JSON.stringify(message)));
+            }
+            catch {
+                // Drop messages that cannot be serialized
+            }
+        }
+        return sanitized;
+    };
     /**
      * The token to use with the secrets manager, setter and getter.
      */

package/lib/{chat-model-registry.d.ts → chat-model-handler.d.ts}

@@ -3,33 +3,30 @@ import { TranslationBundle } from '@jupyterlab/translation';
 import { AgentManagerFactory } from './agent';
 import { AIChatModel } from './chat-model';
 import { AISettingsModel } from './models/settings-model';
-import { IChatModelRegistry, IProviderRegistry, ITokenUsage, IToolRegistry } from './tokens';
+import { IChatModelHandler, IProviderRegistry, ITokenUsage, IToolRegistry } from './tokens';
 import { IDocumentManager } from '@jupyterlab/docmanager';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 /**
- * The chat model registry.
+ * The chat model handler.
  */
-export declare class ChatModelRegistry implements IChatModelRegistry {
-    constructor(options: ChatModelRegistry.IOptions);
-    createModel(name?: string, activeProvider?: string, tokenUsage?: ITokenUsage): AIChatModel;
-    add(model: AIChatModel): void;
-    get(name: string): AIChatModel | undefined;
-    getAll(): AIChatModel[];
-    remove(name: string): void;
+export declare class ChatModelHandler implements IChatModelHandler {
+    constructor(options: ChatModelHandler.IOptions);
+    createModel(name: string, activeProvider: string, tokenUsage?: ITokenUsage): AIChatModel;
     /**
      * Getter/setter for the active cell manager.
      */
     get activeCellManager(): ActiveCellManager | undefined;
     set activeCellManager(manager: ActiveCellManager | undefined);
-    private _models;
     private _docManager;
     private _agentManagerFactory;
     private _settingsModel;
     private _toolRegistry?;
     private _providerRegistry?;
+    private _rmRegistry;
     private _activeCellManager?;
     private _trans;
 }
-export declare namespace ChatModelRegistry {
+export declare namespace ChatModelHandler {
     interface IOptions {
         /**
          * The document manager.
@@ -51,6 +48,10 @@ export declare namespace ChatModelRegistry {
          * Optional provider registry for model creation
          */
         providerRegistry?: IProviderRegistry;
+        /**
+         * Render mime registry.
+         */
+        rmRegistry: IRenderMimeRegistry;
         /**
          * The active cell manager.
          */

package/lib/{chat-model-registry.js → chat-model-handler.js}

@@ -1,15 +1,15 @@
 import { AIChatModel } from './chat-model';
-import { UUID } from '@lumino/coreutils';
 /**
- * The chat model registry.
+ * The chat model handler.
  */
-export class ChatModelRegistry {
+export class ChatModelHandler {
     constructor(options) {
         this._docManager = options.docManager;
         this._agentManagerFactory = options.agentManagerFactory;
         this._settingsModel = options.settingsModel;
         this._toolRegistry = options.toolRegistry;
         this._providerRegistry = options.providerRegistry;
+        this._rmRegistry = options.rmRegistry;
         this._activeCellManager = options.activeCellManager;
         this._trans = options.trans;
     }
@@ -20,7 +20,8 @@ export class ChatModelRegistry {
             toolRegistry: this._toolRegistry,
             providerRegistry: this._providerRegistry,
             activeProvider,
-            tokenUsage
+            tokenUsage,
+            renderMimeRegistry: this._rmRegistry
         });
         // Create AI chat model
         const model = new AIChatModel({
@@ -31,44 +32,9 @@ export class ChatModelRegistry {
             documentManager: this._docManager,
             trans: this._trans
         });
-        // Set the name of the chat if not provided.
-        // The name will be the name set by the user to the model if not already used by
-        // another chat.
-        if (!name || this._models.findIndex(m => m.name === name) !== -1) {
-            const existingName = this.getAll().map(model => model.name);
-            const modelName = this._settingsModel.getProvider(agentManager.activeProvider)?.name ||
-                UUID.uuid4();
-            name = modelName;
-            let i = 1;
-            while (existingName.includes(name)) {
-                name = `${modelName}-${i}`;
-                i += 1;
-            }
-        }
         model.name = name;
-        this.add(model);
         return model;
     }
-    add(model) {
-        if (!this._models.find(m => m.name === model.name)) {
-            this._models.push(model);
-            model.disposed.connect(() => {
-                this.remove(model.name);
-            });
-        }
-    }
-    get(name) {
-        return this._models.find(m => m.name === name);
-    }
-    getAll() {
-        return this._models;
-    }
-    remove(name) {
-        const index = this._models.findIndex(m => m.name === name);
-        if (index !== -1) {
-            this._models.splice(index, 1);
-        }
-    }
     /**
      * Getter/setter for the active cell manager.
      */
@@ -78,12 +44,12 @@ export class ChatModelRegistry {
     set activeCellManager(manager) {
         this._activeCellManager = manager;
     }
-    _models = [];
     _docManager;
     _agentManagerFactory;
     _settingsModel;
     _toolRegistry;
     _providerRegistry;
+    _rmRegistry;
     _activeCellManager;
     _trans;
 }

package/lib/chat-model.d.ts

@@ -92,6 +92,10 @@ export declare class AIChatModel extends AbstractChatModel {
      * @returns A short summary string or empty string if none available
      */
     private _extractToolSummary;
+    /**
+     * Determine whether this tool call should auto-render trusted MIME bundles.
+     */
+    private _computeShouldAutoRenderMimeBundles;
     /**
      * Handles the start of a tool call execution.
      * @param event Event containing the tool call start data
@@ -101,6 +105,10 @@ export declare class AIChatModel extends AbstractChatModel {
      * Handles the completion of a tool call execution.
      */
     private _handleToolCallCompleteEvent;
+    /**
+     * Determine whether a tool call output should auto-render MIME bundles.
+     */
+    private _shouldAutoRenderMimeBundles;
     /**
      * Handles error events from the AI agent.
      */