@jupyterlite/ai 0.13.0 → 0.15.0

package/lib/agent.d.ts

@@ -1,17 +1,16 @@
+import type { IMessageContent } from '@jupyter/chat';
 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';
-type ToolMap = Record<string, Tool>;
+import { type IAgentManager, type IAgentManagerFactory, type IAISettingsModel, type ISkillRegistry, type ITokenUsage, type ToolMap } from './tokens';
+/**
+ * The agent manager factory namespace.
+ */
 export declare namespace AgentManagerFactory {
     interface IOptions {
         /**
          * The settings model.
          */
-        settingsModel: AISettingsModel;
+        settingsModel: IAISettingsModel;
         /**
          * The skill registry for discovering skills.
          */
@@ -26,15 +25,21 @@ export declare namespace AgentManagerFactory {
         token: symbol | null;
     }
 }
-export declare class AgentManagerFactory {
+/**
+ * The agent manager factory.
+ */
+export declare class AgentManagerFactory implements IAgentManagerFactory {
     constructor(options: AgentManagerFactory.IOptions);
-    createAgent(options: IAgentManagerOptions): AgentManager;
+    /**
+     * Create a new agent.
+     */
+    createAgent(options: IAgentManager.IOptions): IAgentManager;
     /**
      * Signal emitted when MCP connection status changes
      */
     get mcpConnectionChanged(): ISignal<this, boolean>;
     /**
-     * Checks if a specific MCP server is connected by server name.
+     * Checks whether a specific MCP server is connected.
      * @param serverName The name of the MCP server to check
      * @returns True if the server is connected, false otherwise
      */
@@ -69,107 +74,22 @@ export declare class AgentManagerFactory {
     private _mcpConnectionChanged;
     private _initQueue;
 }
-/**
- * Event type mapping for type safety with inlined interface definitions
- */
-export interface IAgentEventTypeMap {
-    message_start: {
-        messageId: string;
-    };
-    message_chunk: {
-        messageId: string;
-        chunk: string;
-        fullContent: string;
-    };
-    message_complete: {
-        messageId: string;
-        content: string;
-    };
-    tool_call_start: {
-        callId: string;
-        toolName: string;
-        input: string;
-    };
-    tool_call_complete: {
-        callId: string;
-        toolName: string;
-        outputData: unknown;
-        isError: boolean;
-    };
-    tool_approval_request: {
-        approvalId: string;
-        toolCallId: string;
-        toolName: string;
-        args: unknown;
-    };
-    tool_approval_resolved: {
-        approvalId: string;
-        approved: boolean;
-    };
-    error: {
-        error: Error;
-    };
-}
-/**
- * Events emitted by the AgentManager
- */
-export type IAgentEvent<T extends keyof IAgentEventTypeMap = keyof IAgentEventTypeMap> = T extends keyof IAgentEventTypeMap ? {
-    type: T;
-    data: IAgentEventTypeMap[T];
-} : never;
-/**
- * Configuration options for the AgentManager
- */
-export interface IAgentManagerOptions {
-    /**
-     * AI settings model for configuration
-     */
-    settingsModel: AISettingsModel;
-    /**
-     * Optional tool registry for managing available tools
-     */
-    toolRegistry?: IToolRegistry;
-    /**
-     * Optional provider registry for model creation
-     */
-    providerRegistry?: IProviderRegistry;
-    /**
-     * The skill registry for discovering skills.
-     */
-    skillRegistry?: ISkillRegistry;
-    /**
-     * The secrets manager.
-     */
-    secretsManager?: ISecretsManager;
-    /**
-     * The active provider to use with this agent.
-     */
-    activeProvider?: string;
-    /**
-     * Initial token usage.
-     */
-    tokenUsage?: ITokenUsage;
-    /**
-     * JupyterLab render mime registry for discovering supported MIME types.
-     */
-    renderMimeRegistry?: IRenderMimeRegistry;
-}
 /**
  * Manages the AI agent lifecycle and execution loop.
  * Provides agent initialization, tool management, MCP server integration,
  * and handles the complete agent execution cycle.
  * Emits events for UI updates instead of directly manipulating the chat interface.
  */
-export declare class AgentManager {
+export declare class AgentManager implements IAgentManager {
     /**
      * Creates a new AgentManager instance.
      * @param options Configuration options for the agent manager
      */
-    constructor(options: IAgentManagerOptions);
+    constructor(options: IAgentManager.IOptions);
     /**
      * Signal emitted when agent events occur
      */
-    get agentEvent(): ISignal<this, IAgentEvent>;
+    get agentEvent(): ISignal<this, IAgentManager.IAgentEvent>;
     /**
      * Signal emitted when the active provider has changed.
      */
@@ -200,7 +120,7 @@ export declare class AgentManager {
      * Gets the currently selected tools as a record.
      * @returns Record of selected tools
      */
-    get selectedAgentTools(): Record<string, ITool>;
+    get selectedAgentTools(): ToolMap;
     /**
      * Checks if the current configuration is valid for agent operations.
      * Uses the provider registry to determine if an API key is required.
@@ -210,11 +130,17 @@ export declare class AgentManager {
     /**
      * Clears conversation history and resets agent state.
      */
-    clearHistory(): void;
+    clearHistory(): Promise<void>;
+    /**
+     * Sets the history with a list of messages from the chat.
+     * @param messages The chat messages to set as history
+     */
+    setHistory(messages: IMessageContent[]): void;
     /**
      * Stops the current streaming response by aborting the request.
+     * Resolve any pending approval.
      */
-    stopStreaming(): void;
+    stopStreaming(reason?: string): void;
     /**
      * Approves a pending tool call.
      * @param approvalId The approval ID to approve
@@ -337,5 +263,5 @@ export declare class AgentManager {
     private _initQueue;
     private _agentConfig;
     private _pendingApprovals;
+    private _streaming;
 }
-export {};

package/lib/agent.js

@@ -1,9 +1,13 @@
+import { createMCPClient } from '@ai-sdk/mcp';
+import { PromiseDelegate } from '@lumino/coreutils';
 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';
+/**
+ * The agent manager factory.
+ */
 export class AgentManagerFactory {
     constructor(options) {
         Private.setToken(options.token);
@@ -26,6 +30,9 @@ export class AgentManagerFactory {
             this._secretsManager = undefined;
         }
     }
+    /**
+     * Create a new agent.
+     */
     createAgent(options) {
         const agentManager = new AgentManager({
             ...options,
@@ -33,6 +40,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;
     }
     /**
@@ -42,7 +59,7 @@ export class AgentManagerFactory {
         return this._mcpConnectionChanged;
     }
     /**
-     * Checks if a specific MCP server is connected by server name.
+     * Checks whether a specific MCP server is connected.
      * @param serverName The name of the MCP server to check
      * @returns True if the server is connected, false otherwise
      */
@@ -186,6 +203,7 @@ export class AgentManager {
         this._skills = [];
         this._agentConfig = null;
         this._renderMimeRegistry = options.renderMimeRegistry;
+        this._streaming.resolve();
         this.activeProvider =
             options.activeProvider ?? this._settingsModel.config.defaultProvider;
         // Initialize selected tools to all available tools by default
@@ -291,28 +309,55 @@ export class AgentManager {
     /**
      * Clears conversation history and resets agent state.
      */
-    clearHistory() {
+    async clearHistory() {
         // Stop any ongoing streaming
+        this.stopStreaming('Chat cleared');
+        await this._streaming.promise;
+        // Clear history and token usage
+        this._history = [];
+        this._tokenUsage = { inputTokens: 0, outputTokens: 0 };
+        this._tokenUsageChanged.emit(this._tokenUsage);
+    }
+    /**
+     * Sets the history with a list of messages from the chat.
+     * @param messages The chat messages to set as history
+     */
+    setHistory(messages) {
+        // Stop any ongoing streaming and reject awaiting approvals
         this.stopStreaming();
-        // Reject any pending approvals
         for (const [approvalId, pending] of this._pendingApprovals) {
-            pending.resolve(false, 'Chat cleared');
+            pending.resolve(false, 'Chat history changed');
             this._agentEvent.emit({
                 type: 'tool_approval_resolved',
                 data: { approvalId, approved: false }
             });
         }
         this._pendingApprovals.clear();
-        // Clear history and token usage
-        this._history = [];
-        this._tokenUsage = { inputTokens: 0, outputTokens: 0 };
-        this._tokenUsageChanged.emit(this._tokenUsage);
+        // Convert chat messages to model messages
+        const modelMessages = messages.map(msg => {
+            const isAIMessage = msg.sender.username === 'ai-assistant';
+            return {
+                role: isAIMessage ? 'assistant' : 'user',
+                content: msg.body
+            };
+        });
+        this._history = Private.sanitizeModelMessages(modelMessages);
     }
     /**
      * Stops the current streaming response by aborting the request.
+     * Resolve any pending approval.
      */
-    stopStreaming() {
+    stopStreaming(reason) {
         this._controller?.abort();
+        // Reject any pending approvals
+        for (const [approvalId, pending] of this._pendingApprovals) {
+            pending.resolve(false, reason ?? 'Stream ended by user');
+            this._agentEvent.emit({
+                type: 'tool_approval_resolved',
+                data: { approvalId, approved: false }
+            });
+        }
+        this._pendingApprovals.clear();
     }
     /**
      * Approves a pending tool call.
@@ -352,7 +397,9 @@ export class AgentManager {
      * @param message The user message to respond to (may include processed attachment content)
      */
     async generateResponse(message) {
+        this._streaming = new PromiseDelegate();
         this._controller = new AbortController();
+        const responseHistory = [];
         try {
             // Ensure we have an agent
             if (!this._agent) {
@@ -362,14 +409,14 @@ export class AgentManager {
                 throw new Error('Failed to initialize agent');
             }
             // Add user message to history
-            this._history.push({
+            responseHistory.push({
                 role: 'user',
                 content: message
             });
             let continueLoop = true;
             while (continueLoop) {
                 const result = await this._agent.stream({
-                    messages: this._history,
+                    messages: [...this._history, ...responseHistory],
                     abortSignal: this._controller.signal
                 });
                 const streamResult = await this._processStreamResult(result);
@@ -378,12 +425,12 @@ export class AgentManager {
                 this._updateTokenUsage(await result.usage);
                 // Add response messages to history
                 if (responseMessages.messages?.length) {
-                    this._history.push(...Private.sanitizeModelMessages(responseMessages.messages));
+                    responseHistory.push(...responseMessages.messages);
                 }
                 // Add approval response if processed
                 if (streamResult.approvalResponse) {
                     // Check if the last message is a tool message we can append to
-                    const lastMsg = this._history[this._history.length - 1];
+                    const lastMsg = responseHistory[responseHistory.length - 1];
                     if (lastMsg &&
                         lastMsg.role === 'tool' &&
                         Array.isArray(lastMsg.content) &&
@@ -393,11 +440,13 @@ export class AgentManager {
                     }
                     else {
                         // Add as separate message
-                        this._history.push(streamResult.approvalResponse);
+                        responseHistory.push(streamResult.approvalResponse);
                     }
                 }
                 continueLoop = streamResult.approvalProcessed;
             }
+            // Add the messages to the history only if the response ended without error.
+            this._history.push(...Private.sanitizeModelMessages(responseHistory));
         }
         catch (error) {
             if (error.name !== 'AbortError') {
@@ -409,6 +458,7 @@ export class AgentManager {
         }
         finally {
             this._controller = null;
+            this._streaming.resolve();
         }
     }
     /**
@@ -859,24 +909,102 @@ WEB RETRIEVAL POLICY:
     _initQueue = Promise.resolve();
     _agentConfig;
     _pendingApprovals = new Map();
+    _streaming = new PromiseDelegate();
 }
 var Private;
 (function (Private) {
     /**
-     * Keep only serializable messages by doing a JSON round-trip.
-     * Messages that cannot be serialized are dropped.
+     * Sanitize the messages before adding them to the history.
+     *
+     * 1- Make sure the message sequence is not altered:
+     *   - tool-call messages should have a corresponding tool-result (and vice-versa)
+     *   - tool-approval-request should have a tool-approval-response (and vice-versa)
+     *
+     * 2- 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)));
+            if (message.role === 'assistant') {
+                let newMessage;
+                if (!Array.isArray(message.content)) {
+                    newMessage = message;
+                }
+                else {
+                    // Remove assistant message content without a required response.
+                    const newContent = [];
+                    for (const assistantContent of message.content) {
+                        let isContentValid = true;
+                        if (assistantContent.type === 'tool-call') {
+                            const toolCallId = assistantContent.toolCallId;
+                            isContentValid = !!messages.find(msg => msg.role === 'tool' &&
+                                Array.isArray(msg.content) &&
+                                msg.content.find(content => content.type === 'tool-result' &&
+                                    content.toolCallId === toolCallId));
+                        }
+                        else if (assistantContent.type === 'tool-approval-request') {
+                            const approvalId = assistantContent.approvalId;
+                            isContentValid = !!messages.find(msg => msg.role === 'tool' &&
+                                Array.isArray(msg.content) &&
+                                msg.content.find(content => content.type === 'tool-approval-response' &&
+                                    content.approvalId === approvalId));
+                        }
+                        if (isContentValid) {
+                            newContent.push(assistantContent);
+                        }
+                    }
+                    if (newContent.length) {
+                        newMessage = { ...message, content: newContent };
+                    }
+                }
+                if (newMessage) {
+                    try {
+                        sanitized.push(JSON.parse(JSON.stringify(newMessage)));
+                    }
+                    catch {
+                        // Drop messages that cannot be serialized
+                    }
+                }
+            }
+            else if (message.role === 'tool') {
+                // Remove tool message content without request.
+                const newContent = [];
+                for (const toolContent of message.content) {
+                    let isContentValid = true;
+                    if (toolContent.type === 'tool-result') {
+                        const toolCallId = toolContent.toolCallId;
+                        isContentValid = !!sanitized.find(msg => msg.role === 'assistant' &&
+                            Array.isArray(msg.content) &&
+                            msg.content.find(content => content.type === 'tool-call' &&
+                                content.toolCallId === toolCallId));
+                    }
+                    else if (toolContent.type === 'tool-approval-response') {
+                        const approvalId = toolContent.approvalId;
+                        isContentValid = !!sanitized.find(msg => msg.role === 'assistant' &&
+                            Array.isArray(msg.content) &&
+                            msg.content.find(content => content.type === 'tool-approval-request' &&
+                                content.approvalId === approvalId));
+                    }
+                    if (isContentValid) {
+                        newContent.push(toolContent);
+                    }
+                }
+                if (newContent.length) {
+                    try {
+                        sanitized.push(JSON.parse(JSON.stringify({ ...message, content: newContent })));
+                    }
+                    catch {
+                        // Drop messages that cannot be serialized
+                    }
+                }
             }
-            catch {
-                // Drop messages that cannot be serialized
+            else {
+                // Message is a system or user message.
+                sanitized.push(message);
             }
         }
-        return sanitized;
+        return sanitized.length === messages.length ? sanitized : [];
     };
     /**
      * The token to use with the secrets manager, setter and getter.

package/lib/chat-model-handler.d.ts

@@ -1,17 +1,15 @@
 import { ActiveCellManager } from '@jupyter/chat';
-import { TranslationBundle } from '@jupyterlab/translation';
-import { AgentManagerFactory } from './agent';
-import { AIChatModel } from './chat-model';
-import { AISettingsModel } from './models/settings-model';
-import { IChatModelHandler, IProviderRegistry, ITokenUsage, IToolRegistry } from './tokens';
 import { IDocumentManager } from '@jupyterlab/docmanager';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+import { Contents } from '@jupyterlab/services';
+import { AIChatModel } from './chat-model';
+import type { IAgentManagerFactory, IAISettingsModel, IChatModelHandler, ICreateChatOptions, IProviderRegistry, IToolRegistry } from './tokens';
 /**
  * The chat model handler.
  */
 export declare class ChatModelHandler implements IChatModelHandler {
     constructor(options: ChatModelHandler.IOptions);
-    createModel(name: string, activeProvider: string, tokenUsage?: ITokenUsage): AIChatModel;
+    createModel(options: ICreateChatOptions): AIChatModel;
     /**
      * Getter/setter for the active cell manager.
      */
@@ -24,7 +22,7 @@ export declare class ChatModelHandler implements IChatModelHandler {
     private _providerRegistry?;
     private _rmRegistry;
     private _activeCellManager?;
-    private _trans;
+    private _contentsManager?;
 }
 export declare namespace ChatModelHandler {
     interface IOptions {
@@ -35,11 +33,11 @@ export declare namespace ChatModelHandler {
         /**
          * The agent manager factory.
          */
-        agentManagerFactory: AgentManagerFactory;
+        agentManagerFactory: IAgentManagerFactory;
         /**
          * AI settings model for configuration
          */
-        settingsModel: AISettingsModel;
+        settingsModel: IAISettingsModel;
         /**
          * Optional tool registry for managing available tools
          */
@@ -57,8 +55,8 @@ export declare namespace ChatModelHandler {
          */
         activeCellManager?: ActiveCellManager | undefined;
         /**
-         * The application language translation bundle.
+         * The contents manager.
          */
-        trans: TranslationBundle;
+        contentsManager?: Contents.IManager;
     }
 }

package/lib/chat-model-handler.js

@@ -11,9 +11,10 @@ export class ChatModelHandler {
         this._providerRegistry = options.providerRegistry;
         this._rmRegistry = options.rmRegistry;
         this._activeCellManager = options.activeCellManager;
-        this._trans = options.trans;
+        this._contentsManager = options.contentsManager;
     }
-    createModel(name, activeProvider, tokenUsage) {
+    createModel(options) {
+        const { name, activeProvider, tokenUsage, messages, autosave } = options;
         // Create Agent Manager first so it can be shared
         const agentManager = this._agentManagerFactory.createAgent({
             settingsModel: this._settingsModel,
@@ -30,8 +31,12 @@ export class ChatModelHandler {
             agentManager,
             activeCellManager: this._activeCellManager,
             documentManager: this._docManager,
-            trans: this._trans
+            contentsManager: this._contentsManager
         });
+        messages?.forEach(message => {
+            model.messageAdded({ ...message.content });
+        });
+        model.autosave = autosave ?? false;
         model.name = name;
         return model;
     }
@@ -51,5 +56,5 @@ export class ChatModelHandler {
     _providerRegistry;
     _rmRegistry;
     _activeCellManager;
-    _trans;
+    _contentsManager;
 }