@jupyterlite/ai 0.12.0 → 0.13.0

package/src/chat-model.ts

@@ -3,19 +3,26 @@ import {
   IActiveCellManager,
   IAttachment,
   IChatContext,
+  IMessage,
   IMessageContent,
   INewMessage,
   IUser
 } from '@jupyter/chat';
 
+import { YNotebook } from '@jupyter/ydoc';
+
 import { PathExt } from '@jupyterlab/coreutils';
 
 import { IDocumentManager } from '@jupyterlab/docmanager';
 
 import { IDocumentWidget } from '@jupyterlab/docregistry';
 
+import * as nbformat from '@jupyterlab/nbformat';
+
 import { INotebookModel, Notebook } from '@jupyterlab/notebook';
 
+import { IRenderMime } from '@jupyterlab/rendermime';
+
 import { TranslationBundle } from '@jupyterlab/translation';
 
 import { UUID } from '@lumino/coreutils';
@@ -30,10 +37,6 @@ import { AISettingsModel } from './models/settings-model';
 
 import { ITokenUsage } from './tokens';
 
-import { YNotebook } from '@jupyter/ydoc';
-
-import * as nbformat from '@jupyterlab/nbformat';
-
 /**
  * Tool call status types.
  */
@@ -77,6 +80,10 @@ interface IToolExecutionContext {
    * Human-readable summary extracted from tool input for display.
    */
   summary?: string;
+  /**
+   * Whether this tool call should auto-render trusted MIME bundles on completion.
+   */
+  shouldAutoRenderMimeBundles?: boolean;
 }
 
 /**
@@ -334,8 +341,9 @@ export class AIChatModel extends AbstractChatModel {
       type: 'msg',
       raw_time: false
     };
-    this._currentStreamingMessage = aiMessage;
     this.messageAdded(aiMessage);
+    this._currentStreamingMessage =
+      this.messages.find(message => message.id === aiMessage.id) ?? null;
   }
 
   /**
@@ -347,8 +355,7 @@ export class AIChatModel extends AbstractChatModel {
       this._currentStreamingMessage &&
       this._currentStreamingMessage.id === event.data.messageId
     ) {
-      this._currentStreamingMessage.body = event.data.fullContent;
-      this.messageAdded(this._currentStreamingMessage);
+      this._currentStreamingMessage.update({ body: event.data.fullContent });
     }
   }
 
@@ -361,8 +368,7 @@ export class AIChatModel extends AbstractChatModel {
       this._currentStreamingMessage &&
       this._currentStreamingMessage.id === event.data.messageId
     ) {
-      this._currentStreamingMessage.body = event.data.content;
-      this.messageAdded(this._currentStreamingMessage);
+      this._currentStreamingMessage.update({ body: event.data.content });
       this._currentStreamingMessage = null;
     }
   }
@@ -401,6 +407,21 @@ export class AIChatModel extends AbstractChatModel {
             return parsedInput.name;
           }
           break;
+        case 'browser_fetch':
+          if (parsedInput.url) {
+            return parsedInput.url;
+          }
+          break;
+        case 'web_fetch':
+          if (parsedInput.url) {
+            return parsedInput.url;
+          }
+          break;
+        case 'web_search':
+          if (parsedInput.query) {
+            return `query: "${parsedInput.query}"`;
+          }
+          break;
       }
     } catch {
       // If parsing fails, return empty string
@@ -408,6 +429,30 @@ export class AIChatModel extends AbstractChatModel {
     return '';
   }
 
+  /**
+   * Determine whether this tool call should auto-render trusted MIME bundles.
+   */
+  private _computeShouldAutoRenderMimeBundles(
+    toolName: string,
+    input: string
+  ): boolean {
+    if (toolName !== 'execute_command') {
+      return false;
+    }
+
+    try {
+      const parsedInput = JSON.parse(input);
+      return (
+        typeof parsedInput.commandId === 'string' &&
+        this._settingsModel.config.commandsAutoRenderMimeBundles.includes(
+          parsedInput.commandId
+        )
+      );
+    } catch {
+      return false;
+    }
+  }
+
   /**
    * Handles the start of a tool call execution.
    * @param event Event containing the tool call start data
@@ -420,13 +465,19 @@ export class AIChatModel extends AbstractChatModel {
       event.data.toolName,
       event.data.input
     );
+    const shouldAutoRenderMimeBundles =
+      this._computeShouldAutoRenderMimeBundles(
+        event.data.toolName,
+        event.data.input
+      );
     const context: IToolExecutionContext = {
       toolCallId: event.data.callId,
       messageId,
       toolName: event.data.toolName,
       input: event.data.input,
       status: 'pending',
-      summary
+      summary,
+      shouldAutoRenderMimeBundles
     };
 
     this._toolContexts.set(event.data.callId, context);
@@ -455,11 +506,53 @@ export class AIChatModel extends AbstractChatModel {
   private _handleToolCallCompleteEvent(
     event: IAgentEvent<'tool_call_complete'>
   ): void {
+    const context = this._toolContexts.get(event.data.callId);
     const status = event.data.isError ? 'error' : 'completed';
-    this._updateToolCallUI(event.data.callId, status, event.data.output);
+    this._updateToolCallUI(
+      event.data.callId,
+      status,
+      Private.formatToolOutput(event.data.outputData)
+    );
+
+    if (!event.data.isError && this._shouldAutoRenderMimeBundles(context)) {
+      // Tool results are arbitrary command payloads (often wrapped in
+      // { success, result, outputs, ... }), so extract display outputs
+      // defensively instead of assuming a raw kernel message shape.
+      const mimeBundles = Private.extractMimeBundlesFromUnknown(
+        event.data.outputData,
+        {
+          trustedMimeTypes:
+            this._settingsModel.config.trustedMimeTypesForAutoRender
+        }
+      );
+      for (const bundle of mimeBundles) {
+        this.messageAdded({
+          body: bundle,
+          sender: this._getAIUser(),
+          id: UUID.uuid4(),
+          time: Date.now() / 1000,
+          type: 'msg',
+          raw_time: false
+        });
+      }
+    }
+
     this._toolContexts.delete(event.data.callId);
   }
 
+  /**
+   * Determine whether a tool call output should auto-render MIME bundles.
+   */
+  private _shouldAutoRenderMimeBundles(
+    context: IToolExecutionContext | undefined
+  ): boolean {
+    if (!context) {
+      return false;
+    }
+
+    return !!context.shouldAutoRenderMimeBundles;
+  }
+
   /**
    * Handles error events from the AI agent.
    */
@@ -827,12 +920,119 @@ export class AIChatModel extends AbstractChatModel {
   private _user: IUser;
   private _toolContexts: Map<string, IToolExecutionContext> = new Map();
   private _agentManager: AgentManager;
-  private _currentStreamingMessage: IMessageContent | null = null;
+  private _currentStreamingMessage: IMessage | null = null;
   private _nameChanged = new Signal<AIChatModel, string>(this);
   private _trans: TranslationBundle;
 }
 
 namespace Private {
+  type IMimeBody = Partial<IRenderMime.IMimeModel> &
+    Pick<IRenderMime.IMimeModel, 'data'>;
+  type IDisplayOutput =
+    | nbformat.IDisplayData
+    | nbformat.IDisplayUpdate
+    | nbformat.IExecuteResult;
+
+  const isPlainObject = (value: unknown): value is Record<string, unknown> => {
+    return typeof value === 'object' && value !== null && !Array.isArray(value);
+  };
+
+  const isDisplayOutput = (value: unknown): value is IDisplayOutput => {
+    if (!isPlainObject(value)) {
+      return false;
+    }
+
+    const output = value as nbformat.IOutput;
+    return (
+      nbformat.isDisplayData(output) ||
+      nbformat.isDisplayUpdate(output) ||
+      nbformat.isExecuteResult(output)
+    );
+  };
+
+  const toMimeBundle = (
+    value: IDisplayOutput,
+    trustedMimeTypes: ReadonlySet<string>
+  ): IMimeBody | null => {
+    const data = value.data;
+    if (!isPlainObject(data) || Object.keys(data).length === 0) {
+      return null;
+    }
+
+    return {
+      data: data as IRenderMime.IMimeModel['data'],
+      ...(isPlainObject(value.metadata)
+        ? { metadata: value.metadata as IRenderMime.IMimeModel['metadata'] }
+        : {}),
+      // MIME auto-rendering only runs for explicitly configured command IDs.
+      // Trust handling is configurable to keep risky MIME execution opt-in.
+      ...(Object.keys(data).some(m => trustedMimeTypes.has(m))
+        ? { trusted: true }
+        : {})
+    };
+  };
+
+  /**
+   * Normalize arbitrary tool payloads into canonical display outputs.
+   *
+   * Tool outputs are not guaranteed to be raw Jupyter IOPub messages; they are
+   * often wrapped objects (for example `{ success, result: { outputs: [...] } }`).
+   */
+  const toDisplayOutputs = (value: unknown): IDisplayOutput[] => {
+    if (isDisplayOutput(value)) {
+      return [value];
+    }
+
+    if (Array.isArray(value)) {
+      return value.filter(isDisplayOutput);
+    }
+
+    if (!isPlainObject(value)) {
+      return [];
+    }
+
+    if (Array.isArray(value.outputs)) {
+      return value.outputs.filter(isDisplayOutput);
+    }
+
+    if ('result' in value) {
+      return toDisplayOutputs(value.result);
+    }
+
+    return [];
+  };
+
+  /**
+   * Extract rendermime-ready mime bundles from arbitrary tool results.
+   */
+  export function extractMimeBundlesFromUnknown(
+    content: unknown,
+    options: { trustedMimeTypes?: ReadonlyArray<string> } = {}
+  ): IMimeBody[] {
+    const bundles: IMimeBody[] = [];
+    const outputs = toDisplayOutputs(content);
+    const trustedMimeTypes = new Set(options.trustedMimeTypes ?? []);
+    for (const output of outputs) {
+      const bundle = toMimeBundle(output, trustedMimeTypes);
+      if (bundle) {
+        bundles.push(bundle);
+      }
+    }
+    return bundles;
+  }
+
+  export function formatToolOutput(outputData: unknown): string {
+    if (typeof outputData === 'string') {
+      return outputData;
+    }
+
+    try {
+      return JSON.stringify(outputData, null, 2);
+    } catch {
+      return '[Complex object - cannot serialize]';
+    }
+  }
+
   export function escapeHtml(value: string): string {
     // Prefer the same native escaping approach used in JupyterLab itself
     // (e.g. `@jupyterlab/completer`).
@@ -927,7 +1127,9 @@ namespace Private {
   /**
    * Builds HTML for a tool call display.
    */
-  export function buildToolCallHtml(options: IToolCallHtmlOptions): string {
+  export function buildToolCallHtml(
+    options: IToolCallHtmlOptions
+  ): Partial<IRenderMime.IMimeModel> & Pick<IRenderMime.IMimeModel, 'data'> {
     const { toolName, input, status, summary, output, approvalId, trans } =
       options;
     const config = STATUS_CONFIG[status];
@@ -965,7 +1167,7 @@ namespace Private {
 </div>`;
     }
 
-    return `<details class="jp-ai-tool-call ${config.cssClass}"${openAttr}>
+    const HTMLContent = `<details class="jp-ai-tool-call ${config.cssClass}"${openAttr}>
 <summary class="jp-ai-tool-header">
 <div class="jp-ai-tool-icon">⚡</div>
 <div class="jp-ai-tool-title">${escapedToolName}${summaryHtml}</div>
@@ -974,6 +1176,13 @@ namespace Private {
 <div class="jp-ai-tool-body">${bodyContent}
 </div>
 </details>`;
+
+    return {
+      trusted: true,
+      data: {
+        'text/html': HTMLContent
+      }
+    };
   }
 }
 

package/src/completion/completion-provider.ts

@@ -53,6 +53,11 @@ export class AICompletionProvider implements IInlineCompletionProvider {
       this._updateModel();
     });
     this._updateModel();
+
+    // Disable the secrets manager if the token is empty.
+    if (!options.token) {
+      this._secretsManager = undefined;
+    }
   }
 
   /**
@@ -163,14 +168,23 @@ export class AICompletionProvider implements IInlineCompletionProvider {
 
     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::AICompletionProvider error: the settings manager token is not set.\nYou should disable 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(activeProvider.id);
     }
@@ -316,7 +330,7 @@ export namespace AICompletionProvider {
     /**
      * The token used to request the secrets manager.
      */
-    token: symbol;
+    token: symbol | null;
   }
 }
 
@@ -324,11 +338,11 @@ namespace Private {
   /**
    * 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/components/model-select.tsx

@@ -5,7 +5,6 @@ import { Menu, MenuItem, Typography } from '@mui/material';
 import React, { useCallback, useEffect, useState } from 'react';
 import { AIChatModel } from '../chat-model';
 import { AISettingsModel } from '../models/settings-model';
-
 /**
  * Properties for the model select component.
  */
@@ -199,11 +198,11 @@ export function ModelSelect(props: IModelSelectProps): JSX.Element {
             }}
             sx={{
               backgroundColor: isSelected
-                ? 'var(--jp-brand-color3, rgba(33, 150, 243, 0.1))'
+                ? 'var(--mui-palette-primary-main)'
                 : 'transparent',
               '&:hover': {
                 backgroundColor: isSelected
-                  ? 'var(--jp-brand-color3, rgba(33, 150, 243, 0.15))'
+                  ? 'var(--mui-palette-primary-main)'
                   : 'var(--jp-layout-color1)'
               },
               display: 'flex',
@@ -214,7 +213,7 @@ export function ModelSelect(props: IModelSelectProps): JSX.Element {
             {isSelected ? (
               <CheckIcon
                 sx={{
-                  color: 'var(--jp-brand-color1, #2196F3)',
+                  color: 'var(--jp-ui-inverse-font-color1)',
                   fontSize: 16
                 }}
               />
@@ -227,7 +226,7 @@ export function ModelSelect(props: IModelSelectProps): JSX.Element {
               sx={{
                 fontWeight: isSelected ? 600 : 400,
                 color: isSelected
-                  ? 'var(--jp-brand-color1, #2196F3)'
+                  ? 'var(--jp-ui-inverse-font-color1)'
                   : 'inherit'
               }}
             >

package/src/components/tool-select.tsx

@@ -6,12 +6,14 @@ import BuildIcon from '@mui/icons-material/Build';
 
 import CheckIcon from '@mui/icons-material/Check';
 
-import { Menu, MenuItem, Tooltip, Typography } from '@mui/material';
+import { Divider, Menu, MenuItem, Tooltip, Typography } from '@mui/material';
 
 import React, { useCallback, useEffect, useState } from 'react';
 
-import { INamedTool, IToolRegistry } from '../tokens';
+import { INamedTool, IProviderRegistry, IToolRegistry } from '../tokens';
 import { AIChatModel } from '../chat-model';
+import { AISettingsModel } from '../models/settings-model';
+import { createProviderTools } from '../providers/provider-tools';
 
 const SELECT_ITEM_CLASS = 'jp-AIToolSelect-item';
 
@@ -35,6 +37,16 @@ export interface IToolSelectProps
    */
   onToolSelectionChange: (selectedToolNames: string[]) => void;
 
+  /**
+   * The settings model to compute provider-level web tools.
+   */
+  settingsModel: AISettingsModel;
+
+  /**
+   * Registry for provider metadata used to resolve provider tool capabilities.
+   */
+  providerRegistry: IProviderRegistry;
+
   /**
    * The application language translator.
    */
@@ -49,13 +61,19 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
     toolRegistry,
     onToolSelectionChange,
     toolsEnabled,
+    settingsModel,
+    providerRegistry,
+    model,
     translator: trans
   } = props;
+  const chatContext = model.chatContext as AIChatModel.IAIChatContext;
+  const agentManager = chatContext.agentManager;
 
   const [selectedToolNames, setSelectedToolNames] = useState<string[]>([]);
   const [tools, setTools] = useState<INamedTool[]>(
     toolRegistry?.namedTools || []
   );
+  const [providerToolNames, setProviderToolNames] = useState<string[]>([]);
   const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLElement | null>(null);
   const [menuOpen, setMenuOpen] = useState(false);
 
@@ -103,6 +121,48 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
     }
   }, [toolRegistry]);
 
+  // Track provider-level tools (e.g. web_search/web_fetch).
+  useEffect(() => {
+    if (!agentManager || !toolsEnabled) {
+      setProviderToolNames([]);
+      return;
+    }
+
+    const updateProviderTools = () => {
+      const activeProviderId = agentManager.activeProvider;
+      const providerConfig = settingsModel.getProvider(activeProviderId);
+      if (!providerConfig) {
+        setProviderToolNames([]);
+        return;
+      }
+
+      const providerInfo = providerRegistry.getProviderInfo(
+        providerConfig.provider
+      );
+      const providerTools = createProviderTools({
+        providerInfo,
+        customSettings: providerConfig.customSettings,
+        hasFunctionTools: selectedToolNames.length > 0
+      });
+      setProviderToolNames(Object.keys(providerTools));
+    };
+
+    updateProviderTools();
+    settingsModel.stateChanged.connect(updateProviderTools);
+    agentManager.activeProviderChanged.connect(updateProviderTools);
+
+    return () => {
+      settingsModel.stateChanged.disconnect(updateProviderTools);
+      agentManager.activeProviderChanged.disconnect(updateProviderTools);
+    };
+  }, [
+    settingsModel,
+    providerRegistry,
+    agentManager,
+    selectedToolNames.length,
+    toolsEnabled
+  ]);
+
   // Initialize selected tools to all tools by default
   useEffect(() => {
     if (tools.length > 0 && selectedToolNames.length === 0) {
@@ -113,10 +173,13 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
   }, [tools, selectedToolNames.length, onToolSelectionChange]);
 
   // Don't render if tools are disabled or no tools available
-  if (!toolsEnabled || tools.length === 0) {
+  if (!toolsEnabled || (tools.length === 0 && providerToolNames.length === 0)) {
     return <></>;
   }
 
+  const selectedCount = selectedToolNames.length + providerToolNames.length;
+  const totalCount = tools.length + providerToolNames.length;
+
   return (
     <>
       <TooltippedButton
@@ -125,11 +188,11 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
         }}
         tooltip={trans.__(
           'Tools (%1/%2 selected)',
-          selectedToolNames.length.toString(),
-          tools.length.toString()
+          selectedCount.toString(),
+          totalCount.toString()
         )}
         buttonProps={{
-          ...(selectedToolNames.length === 0 && {
+          ...(selectedCount === 0 && {
             variant: 'outlined'
           }),
           title: trans.__('Select AI Tools'),
@@ -143,7 +206,7 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
           }
         }}
         sx={
-          selectedToolNames.length === 0
+          selectedCount === 0
             ? { backgroundColor: 'var(--jp-layout-color3)' }
             : {}
         }
@@ -198,6 +261,42 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
             </MenuItem>
           </Tooltip>
         ))}
+
+        {providerToolNames.length > 0 && tools.length > 0 && <Divider />}
+
+        {providerToolNames.length > 0 && (
+          <MenuItem disabled>
+            <Typography variant="caption">
+              {trans.__('Provider Tools')}
+            </Typography>
+          </MenuItem>
+        )}
+
+        {providerToolNames.map(toolName => {
+          return (
+            <Tooltip
+              key={toolName}
+              title={trans.__('Enabled via provider settings.')}
+              placement="left"
+            >
+              <MenuItem
+                className={SELECT_ITEM_CLASS}
+                onClick={e => {
+                  // Keep provider-managed tools read-only from this menu.
+                  e.stopPropagation();
+                }}
+              >
+                <CheckIcon
+                  sx={{
+                    marginRight: '8px',
+                    color: 'text.disabled'
+                  }}
+                />
+                <Typography variant="body2">{toolName}</Typography>
+              </MenuItem>
+            </Tooltip>
+          );
+        })}
       </Menu>
     </>
   );
@@ -208,6 +307,8 @@ export function ToolSelect(props: IToolSelectProps): JSX.Element {
  */
 export function createToolSelectItem(
   toolRegistry: IToolRegistry,
+  settingsModel: AISettingsModel,
+  providerRegistry: IProviderRegistry,
   toolsEnabled: boolean = true,
   translator: TranslationBundle
 ): InputToolbarRegistry.IToolbarItem {
@@ -225,6 +326,8 @@ export function createToolSelectItem(
       const toolSelectProps: IToolSelectProps = {
         ...props,
         toolRegistry,
+        settingsModel,
+        providerRegistry,
         onToolSelectionChange,
         toolsEnabled,
         translator