@jupyterlite/ai 0.16.0 → 0.18.0

package/src/chat-model.ts

@@ -5,6 +5,7 @@ import {
   IChatContext,
   IMessage,
   IMessageContent,
+  IMimeModelBody,
   INewMessage,
   IUser
 } from '@jupyter/chat';
@@ -31,9 +32,23 @@ import { Debouncer } from '@lumino/polling';
 
 import { ISignal, Signal } from '@lumino/signaling';
 
+import type { UserContent, ImagePart, FilePart, ModelMessage } from 'ai';
+
 import { AI_AVATAR } from './icons';
 
-import type { IAgentManager, IAISettingsModel, ITokenUsage } from './tokens';
+import type {
+  IAgentManager,
+  IAIChatModel,
+  IAISettingsModel,
+  IProviderRegistry,
+  ITokenUsage
+} from './tokens';
+
+import {
+  modelSupportsAudio,
+  modelSupportsImages,
+  modelSupportsPdf
+} from './providers/model-info';
 
 /**
  * Tool call status types.
@@ -66,10 +81,6 @@ interface IToolExecutionContext {
    * The tool input (formatted).
    */
   input: string;
-  /**
-   * Optional approval ID if awaiting approval.
-   */
-  approvalId?: string;
   /**
    * Current status.
    */
@@ -88,7 +99,7 @@ interface IToolExecutionContext {
  * AI Chat Model implementation that provides chat functionality tool integration,
  * and MCP server support.
  */
-export class AIChatModel extends AbstractChatModel {
+export class AIChatModel extends AbstractChatModel implements IAIChatModel {
   /**
    * Constructs a new AIChatModel instance.
    * @param options Configuration options for the chat model
@@ -106,6 +117,7 @@ export class AIChatModel extends AbstractChatModel {
     this._user = options.user;
     this._agentManager = options.agentManager;
     this._contentsManager = options.contentsManager;
+    this._providerRegistry = options.providerRegistry;
 
     // Listen for agent events
     this._agentManager.agentEvent.connect(this._onAgentEvent, this);
@@ -113,6 +125,13 @@ export class AIChatModel extends AbstractChatModel {
     // Listen for settings changes to update chat behavior
     this._settingsModel.stateChanged.connect(this._onSettingsChanged, this);
 
+    // Rebuild history when the model changes
+    this._agentManager.activeProviderChanged.connect(
+      this._onModelChanged,
+      this
+    );
+    this._settingsModel.stateChanged.connect(this._onModelChanged, this);
+
     this._autosaveDebouncer = new Debouncer(this.save, 3000);
   }
 
@@ -133,6 +152,34 @@ export class AIChatModel extends AbstractChatModel {
     this.setReady();
   }
 
+  /**
+   * A signal emitting when the chat name has changed.
+   */
+  get nameChanged(): ISignal<IAIChatModel, string> {
+    return this._nameChanged;
+  }
+
+  /**
+   * The title of the chat.
+   */
+  get title(): string | null {
+    return this._title;
+  }
+  set title(value: string | null) {
+    this._title = value;
+    if (this.autosave) {
+      this._autosaveDebouncer.invoke();
+    }
+    this._titleChanged.emit(this._title);
+  }
+
+  /**
+   * A signal emitting when the chat title has changed.
+   */
+  get titleChanged(): ISignal<IAIChatModel, string | null> {
+    return this._titleChanged;
+  }
+
   /**
    * Whether to save the chat automatically.
    */
@@ -140,6 +187,9 @@ export class AIChatModel extends AbstractChatModel {
     return this._autosave;
   }
   set autosave(value: boolean) {
+    if (value === this._autosave) {
+      return;
+    }
     this._autosave = value;
     this._autosaveChanged.emit(value);
     if (value) {
@@ -151,7 +201,6 @@ export class AIChatModel extends AbstractChatModel {
         this._autosaveDebouncer.invoke,
         this._autosaveDebouncer
       );
-      this._autosaveDebouncer.invoke();
     } else {
       this.messagesUpdated.disconnect(
         this._autosaveDebouncer.invoke,
@@ -162,22 +211,16 @@ export class AIChatModel extends AbstractChatModel {
         this._autosaveDebouncer
       );
     }
+    this._autosaveDebouncer.invoke();
   }
 
   /**
    * A signal emitting when the autosave flag changed.
    */
-  get autosaveChanged(): ISignal<AIChatModel, boolean> {
+  get autosaveChanged(): ISignal<IAIChatModel, boolean> {
     return this._autosaveChanged;
   }
 
-  /**
-   * A signal emitting when the chat name has changed.
-   */
-  get nameChanged(): ISignal<AIChatModel, string> {
-    return this._nameChanged;
-  }
-
   /**
    * Gets the current user information.
    */
@@ -193,7 +236,7 @@ export class AIChatModel extends AbstractChatModel {
   }
 
   /**
-   * Get the agent manager associated to the model.
+   * The agent manager used in the model.
    */
   get agentManager(): IAgentManager {
     return this._agentManager;
@@ -210,6 +253,7 @@ export class AIChatModel extends AbstractChatModel {
    * Dispose of the model.
    */
   dispose(): void {
+    this.stopStreaming();
     this.messagesUpdated.disconnect(
       this._autosaveDebouncer.invoke,
       this._autosaveDebouncer
@@ -229,7 +273,7 @@ export class AIChatModel extends AbstractChatModel {
       stopStreaming: () => this.stopStreaming(),
       clearMessages: () => this.clearMessages(),
       agentManager: this._agentManager,
-      addSystemMessage: (body: string) => this.addSystemMessage(body)
+      addSystemMessage: (body: string) => this._addSystemMessage(body)
     };
   }
 
@@ -243,16 +287,32 @@ export class AIChatModel extends AbstractChatModel {
   /**
    * Clears all messages from the chat and resets conversation state.
    */
-  clearMessages = (): void => {
+  clearMessages = async (): Promise<void> => {
+    this.stopStreaming();
+    this._messageQueue = [];
+    this._isBusy = false;
+    this._queueMessageId = null;
+    this._currentStreamingMessage = null;
     this.messagesDeleted(0, this.messages.length);
+    this.title = null;
     this._toolContexts.clear();
-    this._agentManager.clearHistory();
+    await this._agentManager.clearHistory();
   };
 
+  /**
+   * Overrides messageAdded to ensure queued messages stay at the bottom.
+   */
+  override messageAdded(message: IMessageContent): void {
+    super.messageAdded(message);
+    if (this._queueMessageId && message.id !== this._queueMessageId) {
+      this._updateQueueUI();
+    }
+  }
+
   /**
    * Adds a non-user message to the chat (used by chat commands).
    */
-  addSystemMessage(body: string): void {
+  private _addSystemMessage(body: string): void {
     const message: IMessageContent = {
       body,
       sender: this._getAIUser(),
@@ -285,10 +345,10 @@ export class AIChatModel extends AbstractChatModel {
       raw_time: false,
       attachments: [...this.input.attachments]
     };
-    this.messageAdded(userMessage);
 
     // Check if we have valid configuration
     if (!this._agentManager.hasValidConfig()) {
+      this.messageAdded(userMessage);
       const errorMessage: IMessageContent = {
         body: 'Please configure your AI settings first. Open the AI Settings to set your API key and model.',
         sender: this._getAIUser(),
@@ -301,27 +361,71 @@ export class AIChatModel extends AbstractChatModel {
       return;
     }
 
+    if (this._isBusy) {
+      this._messageQueue.push({
+        id: UUID.uuid4(),
+        body: message.body,
+        _originalMsg: userMessage
+      });
+      this.input.clearAttachments();
+      this._updateQueueUI();
+      return;
+    }
+
+    this._isBusy = true;
+    this.messageAdded(userMessage);
+    this.input.clearAttachments();
+
+    await this._processMessage(userMessage);
+  }
+
+  /**
+   * Internal method to process attachments and send the message to the agent.
+   */
+  private async _processMessage(userMessage: IMessageContent): Promise<void> {
     try {
-      // Process attachments and add their content to the message
-      let enhancedMessage = message.body;
-      if (this.input.attachments.length > 0) {
-        const attachmentContents = await this._processAttachments(
-          this.input.attachments
+      this.updateWriters([{ user: this._getAIUser() }]);
+
+      let enhancedMessage: UserContent = userMessage.body;
+      if (userMessage.attachments && userMessage.attachments.length > 0) {
+        const providerConfig = this._settingsModel.getProvider(
+          this._agentManager.activeProvider
+        );
+        const supportsImages = modelSupportsImages(
+          providerConfig,
+          this._providerRegistry
+        );
+        const supportsPdf = modelSupportsPdf(
+          providerConfig,
+          this._providerRegistry
+        );
+        const supportsAudio = modelSupportsAudio(
+          providerConfig,
+          this._providerRegistry
         );
-        this.input.clearAttachments();
 
-        if (attachmentContents.length > 0) {
-          enhancedMessage +=
-            '\n\n--- Attached Files ---\n' + attachmentContents.join('\n\n');
-        }
+        enhancedMessage = await Private.processAttachments(
+          userMessage.attachments,
+          this.input.documentManager,
+          userMessage.body,
+          supportsImages,
+          supportsPdf,
+          supportsAudio
+        );
       }
 
-      this.updateWriters([{ user: this._getAIUser() }]);
-
       await this._agentManager.generateResponse(enhancedMessage);
     } catch (error) {
       const errorMessage: IMessageContent = {
-        body: `Error generating AI response: ${(error as Error).message}`,
+        body: '',
+        mime_model: {
+          data: {
+            'application/vnd.jupyter.chat.components': 'error'
+          },
+          metadata: {
+            errorMessage: `Error generating AI response: ${(error as Error).message}`
+          }
+        },
         sender: this._getAIUser(),
         id: UUID.uuid4(),
         time: Date.now() / 1000,
@@ -330,8 +434,102 @@ export class AIChatModel extends AbstractChatModel {
       };
       this.messageAdded(errorMessage);
     } finally {
+      this._drainQueue();
+
+      if (
+        this._settingsModel.config.autoTitle &&
+        (this.messages.length <= 5 || this.title === null)
+      ) {
+        try {
+          this.title = await this.requestTitle();
+        } catch (e) {
+          console.warn('Error while generating a title\n', e);
+        }
+      }
+    }
+  }
+
+  /**
+   * Removes the message-queue chat component.
+   */
+  private _removeQueueUI(): void {
+    if (this._queueMessageId) {
+      const existingMsg = this.messages.find(
+        msg => msg.id === this._queueMessageId
+      );
+      if (existingMsg) {
+        const idx = this.messages.indexOf(existingMsg);
+        if (idx !== -1) {
+          this.messagesDeleted(idx, 1);
+        }
+      }
+      this._queueMessageId = null;
+    }
+  }
+
+  /**
+   * Creates or updates the message-queue chat component.
+   */
+  private _updateQueueUI(): void {
+    this._removeQueueUI();
+
+    if (this._messageQueue.length === 0) {
+      return;
+    }
+
+    const queueBody = {
+      data: {
+        'application/vnd.jupyter.chat.components': 'message-queue'
+      },
+      metadata: {
+        messages: this._messageQueue.map(m => ({
+          id: m.id,
+          body: m.body,
+          attachments: m._originalMsg.attachments
+        })),
+        targetId: this.name
+      }
+    } as IMimeModelBody;
+
+    this._queueMessageId = UUID.uuid4();
+    const queueMessage: IMessageContent = {
+      body: '',
+      mime_model: queueBody,
+      sender: { username: 'system', display_name: '' },
+      id: this._queueMessageId,
+      time: Date.now() / 1000,
+      type: 'msg',
+      raw_time: false
+    };
+    this.messageAdded(queueMessage);
+  }
+
+  /**
+   * Processes the next message in the queue, or marks the agent as idle.
+   */
+  private async _drainQueue(): Promise<void> {
+    if (this._messageQueue.length === 0) {
+      this._isBusy = false;
       this.updateWriters([]);
+      this._removeQueueUI();
+      return;
     }
+
+    // Dequeue and push to chat
+    const next = this._messageQueue.shift()!;
+    next._originalMsg.time = Date.now() / 1000;
+    this.messageAdded(next._originalMsg);
+
+    await this._processMessage(next._originalMsg);
+  }
+
+  /**
+   * Removes a queued message by its ID.
+   * @param messageId The ID of the queued message to remove
+   */
+  removeQueuedMessage(messageId: string): void {
+    this._messageQueue = this._messageQueue.filter(msg => msg.id !== messageId);
+    this._updateQueueUI();
   }
 
   /**
@@ -418,13 +616,39 @@ export class AIChatModel extends AbstractChatModel {
         attachments
       };
     });
-    this.clearMessages();
+    await this.clearMessages();
     this.messagesInserted(0, messages);
-    this._agentManager.setHistory(messages);
+    await this._rebuildHistory();
     this.autosave = content.metadata?.autosave ?? false;
+    this.title = content.metadata?.title ?? null;
     return true;
   };
 
+  /**
+   * Request a title to this chat, regarding the message history.
+   */
+  async requestTitle(): Promise<string> {
+    const history = this.messages
+      .filter(msg => msg.body !== '')
+      .map(
+        msg =>
+          `${msg.sender.username === 'ai-assistant' ? 'assistant' : 'user'}: ${msg.body}`
+      )
+      .join('\n');
+    const messages: ModelMessage[] = [
+      {
+        role: 'system',
+        content:
+          "Generate a concise title (no more than 10 words) for the following conversation. Do not use formatting, quotes, or punctuation. Focus on the subject matter and specific content the user is working on, not on the actions taken (e.g. prefer 'Pandas DataFrame filtering' over 'Opening a notebook'). The title should be a noun phrase describing the topic."
+      },
+      {
+        role: 'user',
+        content: history
+      }
+    ];
+    return this.agentManager.textResponse(messages);
+  }
+
   /**
    * Serialize the model for backup
    */
@@ -436,6 +660,13 @@ export class AIChatModel extends AbstractChatModel {
     const attachmentsList: IAttachment[] = []; // Actual attachments
 
     this.messages.forEach(message => {
+      if (
+        message.content?.mime_model?.data?.[
+          'application/vnd.jupyter.chat.components'
+        ] === 'message-queue'
+      ) {
+        return;
+      }
       let attachmentIndexes: string[] = [];
       if (message.attachments) {
         attachmentIndexes = message.attachments.map(attachment => {
@@ -474,7 +705,8 @@ export class AIChatModel extends AbstractChatModel {
       attachments,
       metadata: {
         provider,
-        autosave: this.autosave
+        autosave: this.autosave,
+        ...(this.title ? { title: this.title } : {})
       }
     };
   }
@@ -500,6 +732,75 @@ export class AIChatModel extends AbstractChatModel {
     // Agent manager handles agent recreation automatically via its own settings listener
   }
 
+  /**
+   * Rebuild history when the active model changes.
+   */
+  private _onModelChanged(): void {
+    const providerConfig = this._settingsModel.getProvider(
+      this._agentManager.activeProvider
+    );
+    const modelKey = providerConfig
+      ? `${providerConfig.provider}:${providerConfig.model}`
+      : undefined;
+    if (modelKey && modelKey !== this._currentModelKey) {
+      this._currentModelKey = modelKey;
+      this._rebuildHistory().catch(e =>
+        console.warn('Failed to rebuild history on model change:', e)
+      );
+    }
+  }
+
+  /**
+   * Rebuilds the agent history from the current messages.
+   * For vision-capable models, re-reads binary attachments from disk.
+   * For text-only models, uses message text only.
+   */
+  private async _rebuildHistory(): Promise<void> {
+    const providerConfig = this._settingsModel.getProvider(
+      this._agentManager.activeProvider
+    );
+    const supportsImages = modelSupportsImages(
+      providerConfig,
+      this._providerRegistry
+    );
+    const supportsPdf = modelSupportsPdf(
+      providerConfig,
+      this._providerRegistry
+    );
+    const supportsAudio = modelSupportsAudio(
+      providerConfig,
+      this._providerRegistry
+    );
+
+    const modelMessages: ModelMessage[] = [];
+    for (const msg of this.messages) {
+      const isAI = msg.sender.username === 'ai-assistant';
+      if (!isAI && msg.attachments?.length) {
+        const enhancedContent = await Private.processAttachments(
+          msg.attachments,
+          this.input.documentManager,
+          msg.body,
+          supportsImages,
+          supportsPdf,
+          supportsAudio
+        );
+
+        modelMessages.push({
+          role: 'user',
+          content: enhancedContent
+        } as ModelMessage);
+      } else if (msg.body) {
+        modelMessages.push({
+          role: isAI ? 'assistant' : 'user',
+          content: msg.body
+        } as ModelMessage);
+      }
+      // Skip messages with empty body like tool calls
+    }
+
+    this._agentManager.setHistory(modelMessages);
+  }
+
   /**
    * Handles events emitted by the agent manager.
    * @param event The event data containing type and payload
@@ -700,13 +1001,18 @@ export class AIChatModel extends AbstractChatModel {
       body: '',
       mime_model: {
         data: {
-          'application/vnd.jupyter.chat.components': 'tool-call'
+          'application/vnd.jupyter.chat.components': 'grouped-tool-calls'
         },
         metadata: {
-          toolName: context.toolName,
-          input: context.input,
-          status: context.status,
-          summary: context.summary
+          toolCalls: [
+            {
+              toolCallId: context.toolCallId,
+              title: `${context.toolName}${context.summary ? ' : ' + context.summary : ''}`,
+              kind: context.toolName,
+              status: 'in_progress',
+              rawInput: context.input
+            }
+          ]
         }
       },
       sender: this._getAIUser(),
@@ -778,7 +1084,15 @@ export class AIChatModel extends AbstractChatModel {
    */
   private _handleErrorEvent(event: IAgentManager.IAgentEvent<'error'>): void {
     this.messageAdded({
-      body: `Error generating response: ${event.data.error.message}`,
+      body: '',
+      mime_model: {
+        data: {
+          'application/vnd.jupyter.chat.components': 'error'
+        },
+        metadata: {
+          errorMessage: `Error generating response: ${event.data.error.message}`
+        }
+      },
       sender: this._getAIUser(),
       id: UUID.uuid4(),
       time: Date.now() / 1000,
@@ -797,7 +1111,6 @@ export class AIChatModel extends AbstractChatModel {
     if (!context) {
       return;
     }
-    context.approvalId = event.data.approvalId;
     context.input = JSON.stringify(event.data.args, null, 2);
     this._updateToolCallUI(event.data.toolCallId, 'awaiting_approval');
   }
@@ -808,15 +1121,13 @@ export class AIChatModel extends AbstractChatModel {
   private _handleToolApprovalResolved(
     event: IAgentManager.IAgentEvent<'tool_approval_resolved'>
   ): void {
-    const context = Array.from(this._toolContexts.values()).find(
-      ctx => ctx.approvalId === event.data.approvalId
-    );
+    const context = this._toolContexts.get(event.data.toolCallId);
     if (!context) {
       return;
     }
 
     const status = event.data.approved ? 'approved' : 'rejected';
-    this._updateToolCallUI(context.toolCallId, status);
+    this._updateToolCallUI(event.data.toolCallId, status);
 
     if (!event.data.approved) {
       this._toolContexts.delete(context.toolCallId);
@@ -847,76 +1158,390 @@ export class AIChatModel extends AbstractChatModel {
     existingMessage.update({
       mime_model: {
         data: {
-          'application/vnd.jupyter.chat.components': 'tool-call'
+          'application/vnd.jupyter.chat.components': 'grouped-tool-calls'
         },
         metadata: {
-          toolName: context.toolName,
-          input: context.input,
-          status: context.status,
-          summary: context.summary,
-          output,
-          targetId: this.name,
-          approvalId: context.approvalId
+          toolCalls: [
+            {
+              toolCallId: context.toolCallId,
+              title: `${context.toolName}${context.summary ? ' : ' + context.summary : ''}`,
+              kind: context.toolName,
+              status: context.status,
+              rawInput: context.input,
+              rawOutput: output,
+              sessionId: this.name,
+              permissionStatus:
+                status === 'awaiting_approval' ? 'pending' : 'resolved',
+              ...(status === 'awaiting_approval' && {
+                permissionOptions: [
+                  { optionId: 'approve', name: 'Approve', kind: 'allow_once' },
+                  { optionId: 'reject', name: 'Reject', kind: 'reject_once' }
+                ]
+              })
+            }
+          ]
         }
       }
     });
   }
 
   /**
-   * Processes file attachments and returns their content as formatted strings.
+   * The current message queue
+   */
+  get messageQueue(): Private.IQueuedItem[] {
+    return this._messageQueue;
+  }
+  set messageQueue(value: Private.IQueuedItem[]) {
+    this._messageQueue = value;
+    this._updateQueueUI();
+    if (this._messageQueue.length > 0 && !this._isBusy) {
+      this._drainQueue();
+    }
+  }
+
+  /**
+   * Whether the chat is busy
+   */
+  get isBusy(): boolean {
+    return this._isBusy;
+  }
+  set isBusy(value: boolean) {
+    this._isBusy = value;
+  }
+
+  // Private fields
+  private _settingsModel: IAISettingsModel;
+  private _user: IUser;
+  private _toolContexts: Map<string, IToolExecutionContext> = new Map();
+  private _agentManager: IAgentManager;
+  private _providerRegistry?: IProviderRegistry;
+  private _currentModelKey: string | undefined;
+  private _currentStreamingMessage: IMessage | null = null;
+  private _nameChanged = new Signal<IAIChatModel, string>(this);
+  private _contentsManager?: Contents.IManager;
+  private _autosave: boolean = false;
+  private _autosaveChanged = new Signal<IAIChatModel, boolean>(this);
+  private _autosaveDebouncer: Debouncer;
+  private _messageQueue: Private.IQueuedItem[] = [];
+  private _isBusy: boolean = false;
+  private _queueMessageId: string | null = null;
+  private _title: string | null = null;
+  private _titleChanged = new Signal<IAIChatModel, string | null>(this);
+}
+
+namespace Private {
+  export interface IQueuedItem {
+    id: string;
+    body: string;
+    _originalMsg: IMessageContent;
+  }
+
+  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>
+  ): IMimeModelBody | 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> } = {}
+  ): IMimeModelBody[] {
+    const bundles: IMimeModelBody[] = [];
+    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]';
+    }
+  }
+
+  /**
+   * Processes file attachments and returns the message content with the attachments.
    * @param attachments Array of file attachments to process
-   * @returns Array of formatted attachment contents
+   * @param documentManager Optional document manager for file operations
+   * @param body The message body
+   * @param supportsImages Whether the model supports images
+   * @param supportsPdf Whether the model supports pdfs
+   * @param supportsAudio Whether the model supports audio
+   * @returns Enhanced message content
    */
-  private async _processAttachments(
-    attachments: IAttachment[]
-  ): Promise<string[]> {
-    const contents: string[] = [];
+  export async function processAttachments(
+    attachments: IAttachment[],
+    documentManager: IDocumentManager | null | undefined,
+    body: string,
+    supportsImages: boolean,
+    supportsPdf: boolean,
+    supportsAudio: boolean
+  ): Promise<UserContent> {
+    const textContents: string[] = [];
+    const includedParts: Array<ImagePart | FilePart> = [];
+    const omittedNames: string[] = [];
+
+    if (!documentManager) {
+      return body;
+    }
 
     for (const attachment of attachments) {
       try {
         if (attachment.type === 'notebook' && attachment.cells?.length) {
-          const cellContents = await this._readNotebookCells(attachment);
+          const cellContents = await readNotebookCells(
+            attachment,
+            documentManager
+          );
           if (cellContents) {
-            contents.push(cellContents);
+            textContents.push(cellContents);
           }
         } else {
-          const fileContent = await this._readFileAttachment(attachment);
-          if (fileContent) {
-            const fileExtension = PathExt.extname(
-              attachment.value
-            ).toLowerCase();
-            const language = fileExtension === '.ipynb' ? 'json' : '';
-            contents.push(
-              `**File: ${attachment.value}**\n\`\`\`${language}\n${fileContent}\n\`\`\``
+          let mimetype = attachment.mimetype;
+          const fileExtension = PathExt.extname(attachment.value).toLowerCase();
+
+          // Fetch mimetype from server metadata if not provided
+          if (!mimetype) {
+            try {
+              const diskModel = await documentManager.services.contents.get(
+                attachment.value,
+                { content: false }
+              );
+              mimetype = diskModel?.mimetype;
+            } catch (e) {
+              console.warn(
+                `Failed to fetch metadata for ${attachment.value}:`,
+                e
+              );
+            }
+          }
+
+          if (mimetype?.startsWith('image/')) {
+            if (supportsImages) {
+              const data = await readBinaryAttachment(
+                attachment,
+                documentManager
+              );
+              if (data) {
+                includedParts.push({
+                  type: 'image',
+                  image: data,
+                  mediaType: mimetype
+                });
+              }
+            } else {
+              omittedNames.push(PathExt.basename(attachment.value));
+            }
+          } else if (mimetype === 'application/pdf') {
+            if (supportsPdf) {
+              const data = await readBinaryAttachment(
+                attachment,
+                documentManager
+              );
+              if (data) {
+                includedParts.push({
+                  type: 'file',
+                  data,
+                  mediaType: mimetype,
+                  filename: PathExt.basename(attachment.value)
+                });
+              }
+            } else {
+              omittedNames.push(PathExt.basename(attachment.value));
+            }
+          } else if (mimetype?.startsWith('audio/')) {
+            if (supportsAudio) {
+              const data = await readBinaryAttachment(
+                attachment,
+                documentManager
+              );
+              if (data) {
+                includedParts.push({
+                  type: 'file',
+                  data,
+                  mediaType: mimetype,
+                  filename: PathExt.basename(attachment.value)
+                });
+              }
+            } else {
+              omittedNames.push(PathExt.basename(attachment.value));
+            }
+          } else {
+            const fileContent = await readFileAttachment(
+              attachment,
+              documentManager
             );
+            if (fileContent) {
+              const language =
+                fileExtension === '.ipynb' ||
+                mimetype === 'application/x-ipynb+json'
+                  ? 'json'
+                  : '';
+              textContents.push(
+                `**File: ${attachment.value}**\n\`\`\`${language}\n${fileContent}\n\`\`\``
+              );
+            }
           }
         }
       } catch (error) {
         console.warn(`Failed to read attachment ${attachment.value}:`, error);
-        contents.push(`**File: ${attachment.value}** (Could not read file)`);
+        textContents.push(
+          `**File: ${attachment.value}** (Could not read file)`
+        );
       }
     }
 
-    return contents;
+    let textPart = body;
+    if (textContents.length > 0) {
+      textPart += '\n\n--- Attached Files ---\n' + textContents.join('\n\n');
+    }
+
+    if (omittedNames.length > 0) {
+      textPart += `\n[Attachments omitted (not supported by this model): ${omittedNames.join(', ')}.]`;
+    }
+
+    return includedParts.length > 0
+      ? [{ type: 'text', text: textPart }, ...includedParts]
+      : textPart;
+  }
+
+  /**
+   * Reads a binary attachment and returns its base64-encoded content.
+   * @param attachment The attachment to read
+   * @param documentManager Optional document manager for file operations
+   * @returns Base64 string or null if unable to read
+   */
+  export async function readBinaryAttachment(
+    attachment: IAttachment,
+    documentManager: IDocumentManager | null | undefined
+  ): Promise<string | null> {
+    if (!documentManager) {
+      return null;
+    }
+
+    try {
+      const diskModel = await documentManager.services.contents.get(
+        attachment.value,
+        { content: true }
+      );
+      if (diskModel?.content && diskModel.format === 'base64') {
+        // Strip whitespace/newlines
+        return (diskModel.content as string).replace(/\s/g, '');
+      }
+      return null;
+    } catch (error) {
+      console.warn(
+        `Failed to read binary attachment ${attachment.value}:`,
+        error
+      );
+      return null;
+    }
   }
 
   /**
    * Reads the content of a notebook cell.
    * @param attachment The notebook attachment to read
+   * @param documentManager Optional document manager for file operations
    * @returns Cell content as string or null if unable to read
    */
-  private async _readNotebookCells(
-    attachment: IAttachment
+  export async function readNotebookCells(
+    attachment: IAttachment,
+    documentManager: IDocumentManager | null | undefined
   ): Promise<string | null> {
-    if (attachment.type !== 'notebook' || !attachment.cells) {
+    if (
+      attachment.type !== 'notebook' ||
+      !attachment.cells ||
+      !documentManager
+    ) {
       return null;
     }
 
     try {
       // Try reading from live notebook if open
-      const widget = this.input.documentManager?.findWidget(
-        attachment.value
-      ) as IDocumentWidget<Notebook, INotebookModel> | undefined;
+      const widget = documentManager.findWidget(attachment.value) as
+        | IDocumentWidget<Notebook, INotebookModel>
+        | undefined;
       let cellData: nbformat.ICell[];
       let kernelLang = 'text';
 
@@ -935,7 +1560,7 @@ export class AIChatModel extends AbstractChatModel {
         kernelLang = String(lang);
       } else {
         // Fallback: reading from disk
-        const model = await this.input.documentManager?.services.contents.get(
+        const model = await documentManager.services.contents.get(
           attachment.value
         );
         if (!model || model.type !== 'notebook') {
@@ -1079,21 +1704,26 @@ export class AIChatModel extends AbstractChatModel {
   /**
    * Reads the content of a file attachment.
    * @param attachment The file attachment to read
+   * @param documentManager Optional document manager for file operations
    * @returns File content as string or null if unable to read
    */
-  private async _readFileAttachment(
-    attachment: IAttachment
+  export async function readFileAttachment(
+    attachment: IAttachment,
+    documentManager: IDocumentManager | null | undefined
   ): Promise<string | null> {
     // Handle both 'file' and 'notebook' types since both have a 'value' path
-    if (attachment.type !== 'file' && attachment.type !== 'notebook') {
+    if (
+      (attachment.type !== 'file' && attachment.type !== 'notebook') ||
+      !documentManager
+    ) {
       return null;
     }
 
     try {
       // Try reading from an open widget first
-      const widget = this.input.documentManager?.findWidget(
-        attachment.value
-      ) as IDocumentWidget<Notebook, INotebookModel> | undefined;
+      const widget = documentManager.findWidget(attachment.value) as
+        | IDocumentWidget<Notebook, INotebookModel>
+        | undefined;
 
       if (widget && widget.context && widget.context.model) {
         const model = widget.context.model;
@@ -1108,7 +1738,7 @@ export class AIChatModel extends AbstractChatModel {
       }
 
       // If not open, load from disk
-      const diskModel = await this.input.documentManager?.services.contents.get(
+      const diskModel = await documentManager.services.contents.get(
         attachment.value
       );
 
@@ -1139,127 +1769,6 @@ export class AIChatModel extends AbstractChatModel {
       return null;
     }
   }
-
-  // Private fields
-  private _settingsModel: IAISettingsModel;
-  private _user: IUser;
-  private _toolContexts: Map<string, IToolExecutionContext> = new Map();
-  private _agentManager: IAgentManager;
-  private _currentStreamingMessage: IMessage | null = null;
-  private _nameChanged = new Signal<AIChatModel, string>(this);
-  private _contentsManager?: Contents.IManager;
-  private _autosave: boolean = false;
-  private _autosaveChanged = new Signal<AIChatModel, boolean>(this);
-  private _autosaveDebouncer: Debouncer;
-}
-
-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]';
-    }
-  }
 }
 
 /**
@@ -1294,6 +1803,10 @@ export namespace AIChatModel {
      * The contents manager.
      */
     contentsManager?: Contents.IManager;
+    /**
+     * Optional provider registry for model capability lookups.
+     */
+    providerRegistry?: IProviderRegistry;
     /**
      * Whether to restore or not the message (default to true)
      */
@@ -1311,7 +1824,7 @@ export namespace AIChatModel {
     /**
      * The clear messages callback.
      */
-    clearMessages: () => void;
+    clearMessages: () => Promise<void>;
     /**
      * Adds an assistant/system message to the chat.
      */
@@ -1350,6 +1863,10 @@ export namespace AIChatModel {
        * Whether the chat is automatically saved.
        */
       autosave?: boolean;
+      /**
+       * An optional title of the chat.
+       */
+      title?: string;
     };
   };
 }