@jupyterlite/ai 0.14.0 → 0.16.0

package/src/agent.ts

@@ -1,3 +1,7 @@
+import { createMCPClient, type MCPClient } from '@ai-sdk/mcp';
+import type { IMessageContent } from '@jupyter/chat';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
+import { PromiseDelegate } from '@lumino/coreutils';
 import { ISignal, Signal } from '@lumino/signaling';
 import {
   ToolLoopAgent,
@@ -5,29 +9,32 @@ import {
   type LanguageModel,
   stepCountIs,
   type StreamTextResult,
-  type Tool,
   type ToolApprovalRequestOutput,
   type TypedToolError,
   type TypedToolOutputDenied,
-  type TypedToolResult
+  type TypedToolResult,
+  type AssistantModelMessage
 } from 'ai';
-import { createMCPClient, type MCPClient } from '@ai-sdk/mcp';
 import { ISecretsManager } from 'jupyter-secrets-manager';
-import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 
-import { AISettingsModel } from './models/settings-model';
 import { createModel } from './providers/models';
+import { getEffectiveContextWindow } from './providers/model-info';
 import {
   createProviderTools,
   type IProviderCustomSettings
 } from './providers/provider-tools';
-import type { IProviderInfo, IProviderRegistry } from './tokens';
 import {
-  ISkillRegistry,
-  ISkillSummary,
-  ITool,
-  IToolRegistry,
-  ITokenUsage,
+  type IAgentManager,
+  type IAgentManagerFactory,
+  type IAISettingsModel,
+  type IProviderInfo,
+  type IProviderRegistry,
+  type ISkillRegistry,
+  type ISkillSummary,
+  type ITool,
+  type IToolRegistry,
+  type ITokenUsage,
+  type ToolMap,
   SECRETS_NAMESPACE
 } from './tokens';
 
@@ -39,8 +46,6 @@ interface IMCPClientWrapper {
   client: MCPClient;
 }
 
-type ToolMap = Record<string, Tool>;
-
 /**
  * Result from processing a stream, including approval info if applicable.
  */
@@ -49,18 +54,25 @@ interface IStreamProcessResult {
    * Whether an approval request was encountered and processed.
    */
   approvalProcessed: boolean;
+  /**
+   * Whether the stream was aborted before completion.
+   */
+  aborted: boolean;
   /**
    * The approval response message to add to history (if approval was processed).
    */
   approvalResponse?: ModelMessage;
 }
 
+/**
+ * The agent manager factory namespace.
+ */
 export namespace AgentManagerFactory {
   export interface IOptions {
     /**
      * The settings model.
      */
-    settingsModel: AISettingsModel;
+    settingsModel: IAISettingsModel;
     /**
      * The skill registry for discovering skills.
      */
@@ -75,7 +87,11 @@ export namespace AgentManagerFactory {
     token: symbol | null;
   }
 }
-export class AgentManagerFactory {
+
+/**
+ * The agent manager factory.
+ */
+export class AgentManagerFactory implements IAgentManagerFactory {
   constructor(options: AgentManagerFactory.IOptions) {
     Private.setToken(options.token);
     this._settingsModel = options.settingsModel;
@@ -104,7 +120,10 @@ export class AgentManagerFactory {
     }
   }
 
-  createAgent(options: IAgentManagerOptions): AgentManager {
+  /**
+   * Create a new agent.
+   */
+  createAgent(options: IAgentManager.IOptions): IAgentManager {
     const agentManager = new AgentManager({
       ...options,
       skillRegistry: this._skillRegistry,
@@ -136,7 +155,7 @@ export class AgentManagerFactory {
   }
 
   /**
-   * 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
    */
@@ -253,8 +272,8 @@ export class AgentManagerFactory {
     });
   }
 
-  private _agentManagers: AgentManager[] = [];
-  private _settingsModel: AISettingsModel;
+  private _agentManagers: IAgentManager[] = [];
+  private _settingsModel: IAISettingsModel;
   private _skillRegistry?: ISkillRegistry;
   private _secretsManager?: ISecretsManager;
   private _mcpClients: IMCPClientWrapper[];
@@ -268,60 +287,6 @@ export class AgentManagerFactory {
 const DEFAULT_TEMPERATURE = 0.7;
 const DEFAULT_MAX_TURNS = 25;
 
-/**
- * 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;
-
 /**
  * Cached configuration used to (re)build the agent.
  */
@@ -335,63 +300,18 @@ interface IAgentConfig {
   shouldUseTools: boolean;
 }
 
-/**
- * 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 class AgentManager {
+export class AgentManager implements IAgentManager {
   /**
    * Creates a new AgentManager instance.
    * @param options Configuration options for the agent manager
    */
-  constructor(options: IAgentManagerOptions) {
+  constructor(options: IAgentManager.IOptions) {
     this._settingsModel = options.settingsModel;
     this._toolRegistry = options.toolRegistry;
     this._providerRegistry = options.providerRegistry;
@@ -402,7 +322,7 @@ export class AgentManager {
     this._history = [];
     this._mcpTools = {};
     this._controller = null;
-    this._agentEvent = new Signal<this, IAgentEvent>(this);
+    this._agentEvent = new Signal<this, IAgentManager.IAgentEvent>(this);
     this._tokenUsage = options.tokenUsage ?? {
       inputTokens: 0,
       outputTokens: 0
@@ -411,6 +331,7 @@ export class AgentManager {
     this._skills = [];
     this._agentConfig = null;
     this._renderMimeRegistry = options.renderMimeRegistry;
+    this._streaming.resolve();
 
     this.activeProvider =
       options.activeProvider ?? this._settingsModel.config.defaultProvider;
@@ -424,7 +345,7 @@ export class AgentManager {
   /**
    * Signal emitted when agent events occur
    */
-  get agentEvent(): ISignal<this, IAgentEvent> {
+  get agentEvent(): ISignal<this, IAgentManager.IAgentEvent> {
     return this._agentEvent;
   }
 
@@ -471,7 +392,17 @@ export class AgentManager {
     return this._activeProvider;
   }
   set activeProvider(value: string) {
+    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);
   }
@@ -491,12 +422,12 @@ export class AgentManager {
    * Gets the currently selected tools as a record.
    * @returns Record of selected tools
    */
-  get selectedAgentTools(): Record<string, ITool> {
+  get selectedAgentTools(): ToolMap {
     if (!this._toolRegistry) {
       return {};
     }
 
-    const result: Record<string, ITool> = {};
+    const result: ToolMap = {};
     for (const name of this._selectedToolNames) {
       const tool: ITool | null = this._toolRegistry.get(name);
       if (tool) {
@@ -539,13 +470,32 @@ export class AgentManager {
   /**
    * Clears conversation history and resets agent state.
    */
-  clearHistory(): void {
+  async clearHistory(): Promise<void> {
     // 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: IMessageContent[]): void {
+    // 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 }
@@ -553,17 +503,33 @@ export class AgentManager {
     }
     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
+      } as ModelMessage;
+    });
+    this._history = Private.sanitizeModelMessages(modelMessages);
   }
 
   /**
    * Stops the current streaming response by aborting the request.
+   * Resolve any pending approval.
    */
-  stopStreaming(): void {
+  stopStreaming(reason?: string): void {
     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();
   }
 
   /**
@@ -606,8 +572,9 @@ export class AgentManager {
    * @param message The user message to respond to (may include processed attachment content)
    */
   async generateResponse(message: string): Promise<void> {
+    this._streaming = new PromiseDelegate();
     this._controller = new AbortController();
-
+    const responseHistory: ModelMessage[] = [];
     try {
       // Ensure we have an agent
       if (!this._agent) {
@@ -619,7 +586,7 @@ export class AgentManager {
       }
 
       // Add user message to history
-      this._history.push({
+      responseHistory.push({
         role: 'user',
         content: message
       });
@@ -627,27 +594,38 @@ export class AgentManager {
       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' &&
@@ -658,12 +636,15 @@ export class AgentManager {
             toolContent.push(...streamResult.approvalResponse.content);
           } 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 as Error).name !== 'AbortError') {
         this._agentEvent.emit({
@@ -671,25 +652,45 @@ export class AgentManager {
           data: { error: error as 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.
    */
   private _updateTokenUsage(
-    usage: { inputTokens?: number; outputTokens?: number } | undefined
+    usage: { inputTokens?: number; outputTokens?: number } | undefined,
+    lastRequestInputTokens?: number
   ): void {
+    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.
+   */
+  private _getActiveContextWindow(): number | undefined {
+    const activeProviderConfig = this._settingsModel.getProvider(
+      this._activeProvider
+    );
+    return getEffectiveContextWindow(
+      activeProviderConfig,
+      this._providerRegistry
+    );
   }
 
   /**
@@ -752,6 +753,13 @@ export class AgentManager {
       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;
@@ -859,7 +867,10 @@ ${richOutputWorkflowInstruction}`;
   ): Promise<IStreamProcessResult> {
     let fullResponse = '';
     let currentMessageId: string | null = null;
-    const processResult: IStreamProcessResult = { approvalProcessed: false };
+    const processResult: IStreamProcessResult = {
+      approvalProcessed: false,
+      aborted: false
+    };
 
     for await (const part of result.fullStream) {
       switch (part.type) {
@@ -921,6 +932,14 @@ ${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;
@@ -1208,98 +1227,8 @@ 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.
-   */
-  private _sanitizeHistory(): void {
-    if (this._history.length === 0) {
-      return;
-    }
-
-    const newHistory: ModelMessage[] = [];
-    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
-   */
-  private _getToolCallIds(message: ModelMessage): string[] {
-    const ids: string[] = [];
-
-    // 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
-   */
-  private _matchesAllToolCalls(
-    message: ModelMessage,
-    callIds: string[]
-  ): boolean {
-    if (message.role !== 'tool' || !Array.isArray(message.content)) {
-      return false;
-    }
-
-    const resultIds = new Set<string>();
-    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
-  private _settingsModel: AISettingsModel;
+  private _settingsModel: IAISettingsModel;
   private _toolRegistry?: IToolRegistry;
   private _providerRegistry?: IProviderRegistry;
   private _skillRegistry?: ISkillRegistry;
@@ -1309,7 +1238,7 @@ WEB RETRIEVAL POLICY:
   private _history: ModelMessage[];
   private _mcpTools: ToolMap;
   private _controller: AbortController | null;
-  private _agentEvent: Signal<this, IAgentEvent>;
+  private _agentEvent: Signal<this, IAgentManager.IAgentEvent>;
   private _tokenUsage: ITokenUsage;
   private _tokenUsageChanged: Signal<this, ITokenUsage>;
   private _activeProvider: string = '';
@@ -1322,25 +1251,123 @@ WEB RETRIEVAL POLICY:
     string,
     { resolve: (approved: boolean, reason?: string) => void }
   > = new Map();
+  private _streaming: PromiseDelegate<void> = new PromiseDelegate();
 }
 
 namespace 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.
    */
   export const sanitizeModelMessages = (
     messages: ModelMessage[]
   ): ModelMessage[] => {
     const sanitized: ModelMessage[] = [];
     for (const message of messages) {
-      try {
-        sanitized.push(JSON.parse(JSON.stringify(message)));
-      } catch {
-        // Drop messages that cannot be serialized
+      if (message.role === 'assistant') {
+        let newMessage: AssistantModelMessage | undefined;
+        if (!Array.isArray(message.content)) {
+          newMessage = message;
+        } else {
+          // Remove assistant message content without a required response.
+          const newContent: typeof message.content = [];
+          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: typeof message.content = [];
+        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 : [];
   };
 
   /**