@jupyterlite/ai 0.14.0 → 0.16.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
@@ -234,9 +160,13 @@ export declare class AgentManager {
      */
     generateResponse(message: string): Promise<void>;
     /**
-     * Updates token usage statistics.
+     * Updates cumulative token usage statistics from a completed model step.
      */
     private _updateTokenUsage;
+    /**
+     * Gets the configured context window for the active provider.
+     */
+    private _getActiveContextWindow;
     /**
      * Initializes the AI agent with current settings and tools.
      * Sets up the agent with model configuration, tools, and MCP tools.
@@ -317,18 +247,6 @@ export declare class AgentManager {
      * 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?;
@@ -349,5 +267,5 @@ export declare class AgentManager {
     private _initQueue;
     private _agentConfig;
     private _pendingApprovals;
+    private _streaming;
 }
-export {};

package/lib/agent.js

@@ -1,9 +1,14 @@
+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 { getEffectiveContextWindow } from './providers/model-info';
 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 +31,9 @@ export class AgentManagerFactory {
             this._secretsManager = undefined;
         }
     }
+    /**
+     * Create a new agent.
+     */
     createAgent(options) {
         const agentManager = new AgentManager({
             ...options,
@@ -52,7 +60,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
      */
@@ -196,6 +204,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
@@ -248,7 +257,14 @@ export class AgentManager {
         return this._activeProvider;
     }
     set activeProvider(value) {
+        const previousProvider = this._activeProvider;
         this._activeProvider = value;
+        // Reset request-level context estimate only when switching between providers.
+        if (previousProvider && previousProvider !== value) {
+            this._tokenUsage.lastRequestInputTokens = undefined;
+        }
+        this._tokenUsage.contextWindow = this._getActiveContextWindow();
+        this._tokenUsageChanged.emit(this._tokenUsage);
         this.initializeAgent();
         this._activeProviderChanged.emit(this._activeProvider);
     }
@@ -301,28 +317,59 @@ 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,
+            contextWindow: this._getActiveContextWindow()
+        };
+        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.
@@ -362,7 +409,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) {
@@ -372,28 +421,39 @@ 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);
-                // Get response messages and update token usage
+                if (streamResult.aborted) {
+                    try {
+                        const responseMessages = await result.response;
+                        if (responseMessages.messages?.length) {
+                            this._history.push(...Private.sanitizeModelMessages(responseMessages.messages));
+                        }
+                    }
+                    catch {
+                        // Aborting before a step finishes leaves no completed response to persist.
+                    }
+                    break;
+                }
+                // Get response messages for completed steps.
                 const responseMessages = await result.response;
-                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) &&
@@ -403,11 +463,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') {
@@ -416,23 +478,32 @@ 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;
+            this._streaming.resolve();
         }
     }
     /**
-     * Updates token usage statistics.
+     * Updates cumulative token usage statistics from a completed model step.
      */
-    _updateTokenUsage(usage) {
+    _updateTokenUsage(usage, lastRequestInputTokens) {
+        const contextWindow = this._getActiveContextWindow();
+        const estimatedRequestInputTokens = lastRequestInputTokens ?? usage?.inputTokens;
         if (usage) {
             this._tokenUsage.inputTokens += usage.inputTokens ?? 0;
             this._tokenUsage.outputTokens += usage.outputTokens ?? 0;
-            this._tokenUsageChanged.emit(this._tokenUsage);
         }
+        this._tokenUsage.lastRequestInputTokens = estimatedRequestInputTokens;
+        this._tokenUsage.contextWindow = contextWindow;
+        this._tokenUsageChanged.emit(this._tokenUsage);
+    }
+    /**
+     * Gets the configured context window for the active provider.
+     */
+    _getActiveContextWindow() {
+        const activeProviderConfig = this._settingsModel.getProvider(this._activeProvider);
+        return getEffectiveContextWindow(activeProviderConfig, this._providerRegistry);
     }
     /**
      * Initializes the AI agent with current settings and tools.
@@ -484,6 +555,9 @@ export class AgentManager {
         const activeProviderInfo = activeProviderConfig && this._providerRegistry
             ? this._providerRegistry.getProviderInfo(activeProviderConfig.provider)
             : null;
+        const contextWindow = getEffectiveContextWindow(activeProviderConfig, this._providerRegistry);
+        this._tokenUsage.contextWindow = contextWindow;
+        this._tokenUsageChanged.emit(this._tokenUsage);
         const temperature = activeProviderConfig?.parameters?.temperature ?? DEFAULT_TEMPERATURE;
         const maxTokens = activeProviderConfig?.parameters?.maxOutputTokens;
         const maxTurns = activeProviderConfig?.parameters?.maxTurns ?? DEFAULT_MAX_TURNS;
@@ -562,7 +636,10 @@ ${richOutputWorkflowInstruction}`;
     async _processStreamResult(result) {
         let fullResponse = '';
         let currentMessageId = null;
-        const processResult = { approvalProcessed: false };
+        const processResult = {
+            approvalProcessed: false,
+            aborted: false
+        };
         for await (const part of result.fullStream) {
             switch (part.type) {
                 case 'text-delta':
@@ -617,6 +694,12 @@ ${richOutputWorkflowInstruction}`;
                     }
                     await this._handleApprovalRequest(part, processResult);
                     break;
+                case 'finish-step':
+                    this._updateTokenUsage(part.usage, part.usage.inputTokens);
+                    break;
+                case 'abort':
+                    processResult.aborted = true;
+                    break;
                 // Ignore: text-start, text-end, finish, error, and others
                 default:
                     break;
@@ -851,81 +934,6 @@ WEB RETRIEVAL POLICY:
         }
         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;
     _toolRegistry;
@@ -947,24 +955,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
+                    }
+                }
             }
-            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
+                    }
+                }
+            }
+            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.