@jupyterlite/ai 0.12.0 → 0.14.0

package/src/agent.ts

@@ -7,14 +7,21 @@ import {
   type StreamTextResult,
   type Tool,
   type ToolApprovalRequestOutput,
+  type TypedToolError,
+  type TypedToolOutputDenied,
   type TypedToolResult
 } 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 type { IProviderRegistry } from './tokens';
+import {
+  createProviderTools,
+  type IProviderCustomSettings
+} from './providers/provider-tools';
+import type { IProviderInfo, IProviderRegistry } from './tokens';
 import {
   ISkillRegistry,
   ISkillSummary,
@@ -65,7 +72,7 @@ export namespace AgentManagerFactory {
     /**
      * The token used to request the secrets manager.
      */
-    token: symbol;
+    token: symbol | null;
   }
 }
 export class AgentManagerFactory {
@@ -90,6 +97,11 @@ export class AgentManagerFactory {
 
     // 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: IAgentManagerOptions): AgentManager {
@@ -99,6 +111,20 @@ 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;
   }
 
@@ -266,7 +292,7 @@ export interface IAgentEventTypeMap {
   tool_call_complete: {
     callId: string;
     toolName: string;
-    output: string;
+    outputData: unknown;
     isError: boolean;
   };
   tool_approval_request: {
@@ -347,6 +373,11 @@ export interface IAgentManagerOptions {
    * Initial token usage.
    */
   tokenUsage?: ITokenUsage;
+
+  /**
+   * JupyterLab render mime registry for discovering supported MIME types.
+   */
+  renderMimeRegistry?: IRenderMimeRegistry;
 }
 
 /**
@@ -379,6 +410,7 @@ export class AgentManager {
     this._tokenUsageChanged = new Signal<this, ITokenUsage>(this);
     this._skills = [];
     this._agentConfig = null;
+    this._renderMimeRegistry = options.renderMimeRegistry;
 
     this.activeProvider =
       options.activeProvider ?? this._settingsModel.config.defaultProvider;
@@ -607,7 +639,9 @@ export class AgentManager {
 
         // 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
@@ -637,6 +671,9 @@ 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;
     }
@@ -697,21 +734,24 @@ export class AgentManager {
 
     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 supportsToolCalling = this._supportsToolCalling();
+    const canUseTools = config.toolsEnabled && supportsToolCalling;
+    const hasFunctionToolRegistry = !!(
+      this._toolRegistry && Object.keys(this._toolRegistry.tools).length > 0
     );
-
-    const tools = shouldUseTools
-      ? { ...this.selectedAgentTools, ...this._mcpTools }
-      : this._mcpTools;
+    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;
@@ -719,6 +759,15 @@ export class AgentManager {
     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,
@@ -730,6 +779,29 @@ export class AgentManager {
     };
   }
 
+  /**
+   * Build the runtime tool map used by the agent.
+   */
+  private _buildRuntimeTools(options: {
+    providerInfo?: IProviderInfo | null;
+    customSettings?: IProviderCustomSettings;
+    functionTools: ToolMap;
+    includeProviderTools: boolean;
+  }): ToolMap {
+    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.
    */
@@ -749,9 +821,22 @@ export class AgentManager {
       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,
@@ -818,6 +903,14 @@ export class AgentManager {
           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) {
@@ -856,10 +949,6 @@ export class AgentManager {
    * Handles tool-result stream parts.
    */
   private _handleToolResult(part: TypedToolResult<ToolMap>): void {
-    const output =
-      typeof part.output === 'string'
-        ? part.output
-        : JSON.stringify(part.output, null, 2);
     const isError =
       typeof part.output === 'object' &&
       part.output !== null &&
@@ -871,12 +960,49 @@ export class AgentManager {
       data: {
         callId: part.toolCallId,
         toolName: part.toolName,
-        output,
+        outputData: part.output,
         isError
       }
     });
   }
 
+  /**
+   * Handles tool-error stream parts.
+   */
+  private _handleToolError(part: TypedToolError<ToolMap>): void {
+    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.
+   */
+  private _handleToolOutputDenied(part: TypedToolOutputDenied<ToolMap>): void {
+    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.
    */
@@ -980,14 +1106,23 @@ export class AgentManager {
 
     let apiKey: string;
     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);
     }
@@ -1008,15 +1143,17 @@ export class AgentManager {
    * @param baseSystemPrompt The base system prompt from settings
    * @returns The enhanced system prompt with dynamic additions
    */
-  private _getEnhancedSystemPrompt(baseSystemPrompt: string): string {
-    if (this._skills.length === 0) {
-      return baseSystemPrompt;
-    }
-
-    const lines = this._skills.map(
-      skill => `- ${skill.name}: ${skill.description}`
-    );
-    const skillsPrompt = `
+  private _getEnhancedSystemPrompt(
+    baseSystemPrompt: string,
+    tools: ToolMap
+  ): string {
+    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).
@@ -1028,8 +1165,137 @@ If the load_skill result includes a non-empty "resources" array, those are bundl
 AVAILABLE SKILLS (preloaded snapshot):
 ${lines.join('\n')}
 `;
+      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.
+   */
+  private _getSupportedMimeTypesInstruction(): string {
+    const mimeTypes = this._renderMimeRegistry?.mimeTypes ?? [];
+    const safeMimeTypes = mimeTypes.filter(mimeType => {
+      const factory = this._renderMimeRegistry?.getFactory(mimeType);
+      return !!factory?.safe;
+    });
 
-    return baseSystemPrompt + skillsPrompt;
+    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.
+   */
+  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
@@ -1049,6 +1315,7 @@ ${lines.join('\n')}
   private _activeProvider: string = '';
   private _activeProviderChanged = new Signal<this, string | undefined>(this);
   private _skills: ISkillSummary[];
+  private _renderMimeRegistry?: IRenderMimeRegistry;
   private _initQueue: Promise<void> = Promise.resolve();
   private _agentConfig: IAgentConfig | null;
   private _pendingApprovals: Map<
@@ -1058,14 +1325,32 @@ ${lines.join('\n')}
 }
 
 namespace Private {
+  /**
+   * 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
+      }
+    }
+    return sanitized;
+  };
+
   /**
    * The token to use with the secrets manager, setter and getter.
    */
-  let secretsToken: symbol;
-  export function setToken(value: symbol): void {
+  let secretsToken: symbol | null;
+  export function setToken(value: symbol | null): void {
     secretsToken = value;
   }
-  export function getToken(): symbol {
+  export function getToken(): symbol | null {
     return secretsToken;
   }
 }

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

@@ -4,31 +4,32 @@ import { AgentManagerFactory } from './agent';
 import { AIChatModel } from './chat-model';
 import { AISettingsModel } from './models/settings-model';
 import {
-  IChatModelRegistry,
+  IChatModelHandler,
   IProviderRegistry,
   ITokenUsage,
   IToolRegistry
 } from './tokens';
 import { IDocumentManager } from '@jupyterlab/docmanager';
-import { UUID } from '@lumino/coreutils';
+import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 
 /**
- * The chat model registry.
+ * The chat model handler.
  */
-export class ChatModelRegistry implements IChatModelRegistry {
-  constructor(options: ChatModelRegistry.IOptions) {
+export class ChatModelHandler implements IChatModelHandler {
+  constructor(options: ChatModelHandler.IOptions) {
     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;
   }
 
   createModel(
-    name?: string,
-    activeProvider?: string,
+    name: string,
+    activeProvider: string,
     tokenUsage?: ITokenUsage
   ): AIChatModel {
     // Create Agent Manager first so it can be shared
@@ -37,7 +38,8 @@ export class ChatModelRegistry implements IChatModelRegistry {
       toolRegistry: this._toolRegistry,
       providerRegistry: this._providerRegistry,
       activeProvider,
-      tokenUsage
+      tokenUsage,
+      renderMimeRegistry: this._rmRegistry
     });
 
     // Create AI chat model
@@ -50,52 +52,11 @@ export class ChatModelRegistry implements IChatModelRegistry {
       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: AIChatModel): void {
-    if (!this._models.find(m => m.name === model.name)) {
-      this._models.push(model);
-      model.disposed.connect(() => {
-        this.remove(model.name);
-      });
-    }
-  }
-
-  get(name: string): AIChatModel | undefined {
-    return this._models.find(m => m.name === name);
-  }
-
-  getAll(): AIChatModel[] {
-    return this._models;
-  }
-
-  remove(name: string): void {
-    const index = this._models.findIndex(m => m.name === name);
-    if (index !== -1) {
-      this._models.splice(index, 1);
-    }
-  }
-
   /**
    * Getter/setter for the active cell manager.
    */
@@ -106,17 +67,17 @@ export class ChatModelRegistry implements IChatModelRegistry {
     this._activeCellManager = manager;
   }
 
-  private _models: AIChatModel[] = [];
   private _docManager: IDocumentManager;
   private _agentManagerFactory: AgentManagerFactory;
   private _settingsModel: AISettingsModel;
   private _toolRegistry?: IToolRegistry;
   private _providerRegistry?: IProviderRegistry;
+  private _rmRegistry: IRenderMimeRegistry;
   private _activeCellManager?: ActiveCellManager;
   private _trans: TranslationBundle;
 }
 
-export namespace ChatModelRegistry {
+export namespace ChatModelHandler {
   export interface IOptions {
     /**
      * The document manager.
@@ -138,6 +99,10 @@ export namespace ChatModelRegistry {
      * Optional provider registry for model creation
      */
     providerRegistry?: IProviderRegistry;
+    /**
+     * Render mime registry.
+     */
+    rmRegistry: IRenderMimeRegistry;
     /**
      * The active cell manager.
      */