diff --git a/docs/usage.md b/docs/usage.md index 8d9557a6..954b660f 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -32,3 +32,22 @@ See the dedicated pages for specific providers: - [Using Ollama](./ollama.md) - [Using LiteLLM Proxy](./litellm.md) - [Using any-llm-gateway](./any-llm-gateway.md) + +## Controlling MIME auto-rendering in chat + +When the AI model uses `execute_command`, some commands may return rich MIME bundles +(plots, maps, HTML, etc.). You can control which commands automatically render +those bundles as chat messages: + +1. Open AI settings and go to **Behavior Settings** +2. In **Commands Auto-Rendering MIME Bundles**, add or remove command IDs +3. In **Trusted MIME Types for Auto-Render**, add or remove MIME types to mark + as trusted when those commands are auto-rendered in chat + +Default: + +- `jupyterlab-ai-commands:execute-in-kernel` +- `text/html` (trusted MIME type) + +This helps avoid side effects where inspection commands return existing notebook +outputs that you do not want replayed in chat. diff --git a/schema/settings-model.json b/schema/settings-model.json index 822083ee..957d19e8 100644 --- a/schema/settings-model.json +++ b/schema/settings-model.json @@ -159,6 +159,20 @@ "jupyterlab-ai-commands:run-cell" ] }, + "commandsAutoRenderMimeBundles": { + "title": "Commands Auto-Rendering MIME Bundles", + "description": "List of execute_command command IDs whose outputs can auto-render MIME bundles in chat", + "type": "array", + "items": { "type": "string" }, + "default": ["jupyterlab-ai-commands:execute-in-kernel"] + }, + "trustedMimeTypesForAutoRender": { + "title": "Trusted MIME Types for Auto-Render", + "description": "MIME types to mark as trusted when auto-rendering execute_command outputs in chat", + "type": "array", + "items": { "type": "string" }, + "default": ["text/html"] + }, "systemPrompt": { "title": "System Prompt", "description": "Instructions that define how the AI should behave and respond", diff --git a/src/agent.ts b/src/agent.ts index 8a484c4e..41b7d7f5 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -11,6 +11,7 @@ import { } 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'; @@ -271,7 +272,7 @@ export interface IAgentEventTypeMap { tool_call_complete: { callId: string; toolName: string; - output: string; + outputData: unknown; isError: boolean; }; tool_approval_request: { @@ -352,6 +353,11 @@ export interface IAgentManagerOptions { * Initial token usage. */ tokenUsage?: ITokenUsage; + + /** + * JupyterLab render mime registry for discovering supported MIME types. + */ + renderMimeRegistry?: IRenderMimeRegistry; } /** @@ -384,6 +390,7 @@ export class AgentManager { this._tokenUsageChanged = new Signal(this); this._skills = []; this._agentConfig = null; + this._renderMimeRegistry = options.renderMimeRegistry; this.activeProvider = options.activeProvider ?? this._settingsModel.config.defaultProvider; @@ -612,7 +619,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 @@ -754,9 +763,22 @@ export class AgentManager { shouldUseTools } = this._agentConfig; - const instructions = shouldUseTools + const baseInstructions = shouldUseTools ? this._getEnhancedSystemPrompt(baseSystemPrompt) : 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, @@ -861,10 +883,6 @@ export class AgentManager { * Handles tool-result stream parts. */ private _handleToolResult(part: TypedToolResult): void { - const output = - typeof part.output === 'string' - ? part.output - : JSON.stringify(part.output, null, 2); const isError = typeof part.output === 'object' && part.output !== null && @@ -876,7 +894,7 @@ export class AgentManager { data: { callId: part.toolCallId, toolName: part.toolName, - output, + outputData: part.output, isError } }); @@ -1046,6 +1064,23 @@ ${lines.join('\n')} return baseSystemPrompt + skillsPrompt; } + /** + * 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; + }); + + 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(', ')}`; + } + // Private attributes private _settingsModel: AISettingsModel; private _toolRegistry?: IToolRegistry; @@ -1063,6 +1098,7 @@ ${lines.join('\n')} private _activeProvider: string = ''; private _activeProviderChanged = new Signal(this); private _skills: ISkillSummary[]; + private _renderMimeRegistry?: IRenderMimeRegistry; private _initQueue: Promise = Promise.resolve(); private _agentConfig: IAgentConfig | null; private _pendingApprovals: Map< @@ -1072,6 +1108,24 @@ ${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. */ diff --git a/src/chat-model-registry.ts b/src/chat-model-registry.ts index db80a52e..7e3e3b5e 100644 --- a/src/chat-model-registry.ts +++ b/src/chat-model-registry.ts @@ -11,6 +11,7 @@ import { } from './tokens'; import { IDocumentManager } from '@jupyterlab/docmanager'; import { UUID } from '@lumino/coreutils'; +import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; /** * The chat model registry. @@ -22,6 +23,7 @@ export class ChatModelRegistry implements IChatModelRegistry { this._settingsModel = options.settingsModel; this._toolRegistry = options.toolRegistry; this._providerRegistry = options.providerRegistry; + this._rmRegistry = options.rmRegistry; this._activeCellManager = options.activeCellManager; this._trans = options.trans; } @@ -37,7 +39,8 @@ export class ChatModelRegistry implements IChatModelRegistry { toolRegistry: this._toolRegistry, providerRegistry: this._providerRegistry, activeProvider, - tokenUsage + tokenUsage, + renderMimeRegistry: this._rmRegistry }); // Create AI chat model @@ -112,6 +115,7 @@ export class ChatModelRegistry implements IChatModelRegistry { private _settingsModel: AISettingsModel; private _toolRegistry?: IToolRegistry; private _providerRegistry?: IProviderRegistry; + private _rmRegistry: IRenderMimeRegistry; private _activeCellManager?: ActiveCellManager; private _trans: TranslationBundle; } @@ -138,6 +142,10 @@ export namespace ChatModelRegistry { * Optional provider registry for model creation */ providerRegistry?: IProviderRegistry; + /** + * Render mime registry. + */ + rmRegistry: IRenderMimeRegistry; /** * The active cell manager. */ diff --git a/src/chat-model.ts b/src/chat-model.ts index eedc7b2d..1b629ada 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -80,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; } /** @@ -410,6 +414,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 @@ -422,13 +450,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); @@ -457,11 +491,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. */ @@ -835,6 +911,113 @@ export class AIChatModel extends AbstractChatModel { } namespace Private { + type IMimeBody = Partial & + Pick; + type IDisplayOutput = + | nbformat.IDisplayData + | nbformat.IDisplayUpdate + | nbformat.IExecuteResult; + + const isPlainObject = (value: unknown): value is Record => { + 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 + ): 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 } = {} + ): 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`). @@ -980,8 +1163,8 @@ namespace Private { `; return { + trusted: true, data: { - trusted: true, 'text/html': HTMLContent } }; diff --git a/src/index.ts b/src/index.ts index 37ede728..c9e07379 100644 --- a/src/index.ts +++ b/src/index.ts @@ -58,6 +58,7 @@ import { DisposableSet } from '@lumino/disposable'; import { AgentManagerFactory } from './agent'; import { AIChatModel } from './chat-model'; +import { RenderedMessageOutputAreaCompat } from './rendered-message-outputarea'; import { ClearCommandProvider } from './chat-commands/clear'; import { SkillsCommandProvider } from './chat-commands/skills'; @@ -250,7 +251,12 @@ const chatModelRegistry: JupyterFrontEndPlugin = { id: '@jupyterlite/ai:chat-model-registry', description: 'Registry for the current chat model', autoStart: true, - requires: [IAISettingsModel, IAgentManagerFactory, IDocumentManager], + requires: [ + IAISettingsModel, + IAgentManagerFactory, + IDocumentManager, + IRenderMimeRegistry + ], optional: [IProviderRegistry, IToolRegistry, ITranslator], provides: IChatModelRegistry, activate: ( @@ -258,6 +264,7 @@ const chatModelRegistry: JupyterFrontEndPlugin = { settingsModel: AISettingsModel, agentManagerFactory: AgentManagerFactory, docManager: IDocumentManager, + rmRegistry: IRenderMimeRegistry, providerRegistry?: IProviderRegistry, toolRegistry?: IToolRegistry, translator?: ITranslator @@ -268,6 +275,7 @@ const chatModelRegistry: JupyterFrontEndPlugin = { settingsModel, agentManagerFactory, docManager, + rmRegistry, providerRegistry, toolRegistry, trans @@ -415,10 +423,16 @@ const plugin: JupyterFrontEndPlugin = { chatPanel: widget, agentManager: model.agentManager }); + // Temporary compat: keep output-area CSS context for MIME renderers + // until jupyter-chat provides it natively. + const outputAreaCompat = new RenderedMessageOutputAreaCompat({ + chatPanel: widget + }); widget.disposed.connect(() => { // Dispose of the approval buttons widget when the chat is disposed. approvalButton.dispose(); + outputAreaCompat.dispose(); // Remove the model from the registry when the widget is disposed. modelRegistry.remove(model.name); }); diff --git a/src/models/settings-model.ts b/src/models/settings-model.ts index a8357dea..15f713ff 100644 --- a/src/models/settings-model.ts +++ b/src/models/settings-model.ts @@ -56,6 +56,10 @@ export interface IAIConfig { showTokenUsage: boolean; // Commands that require approval before execution commandsRequiringApproval: string[]; + // Commands whose execute_command outputs may auto-render MIME bundles in chat + commandsAutoRenderMimeBundles: string[]; + // MIME types that are trusted when auto-rendering execute_command outputs + trustedMimeTypesForAutoRender: string[]; // Diff display settings showCellDiff: boolean; showFileDiff: boolean; @@ -97,6 +101,8 @@ export class AISettingsModel extends VDomModel { 'runmenu:run-all', 'jupyterlab-ai-commands:run-cell' ], + commandsAutoRenderMimeBundles: ['jupyterlab-ai-commands:execute-in-kernel'], + trustedMimeTypesForAutoRender: ['text/html'], systemPrompt: `You are Jupyternaut, an AI coding assistant built specifically for the JupyterLab environment. ## Your Core Mission diff --git a/src/rendered-message-outputarea.ts b/src/rendered-message-outputarea.ts new file mode 100644 index 00000000..2f0b9eaf --- /dev/null +++ b/src/rendered-message-outputarea.ts @@ -0,0 +1,62 @@ +import { ChatWidget } from '@jupyter/chat'; +import { IDisposable } from '@lumino/disposable'; + +const OUTPUT_AREA_CLASS = 'jp-OutputArea'; +const CHAT_RENDERED_MESSAGE_SELECTOR = `.jp-chat-rendered-message:not(.${OUTPUT_AREA_CLASS})`; + +/** + * Ensures chat-rendered MIME outputs also expose the OutputArea class so + * renderer extensions can reuse their notebook/output-area CSS rules. + * + * TODO: Remove this compatibility layer once jupyter-chat applies + * `jp-OutputArea` (or equivalent output-area context) to rendered MIME + * messages by default. + */ +export class RenderedMessageOutputAreaCompat implements IDisposable { + constructor(options: RenderedMessageOutputAreaCompat.IOptions) { + this._chatPanel = options.chatPanel; + this._chatPanel.model.messagesUpdated.connect(this._scheduleSync, this); + this._scheduleSync(); + } + + get isDisposed(): boolean { + return this._isDisposed; + } + + dispose(): void { + if (this._isDisposed) { + return; + } + this._isDisposed = true; + this._chatPanel.model.messagesUpdated.disconnect(this._scheduleSync, this); + if (this._raf !== 0) { + cancelAnimationFrame(this._raf); + this._raf = 0; + } + } + + private _scheduleSync(): void { + if (this._isDisposed || this._raf !== 0) { + return; + } + this._raf = requestAnimationFrame(() => { + this._raf = 0; + if (this._isDisposed) { + return; + } + this._chatPanel.node + .querySelectorAll(CHAT_RENDERED_MESSAGE_SELECTOR) + .forEach(element => element.classList.add(OUTPUT_AREA_CLASS)); + }); + } + + private readonly _chatPanel: ChatWidget; + private _isDisposed = false; + private _raf = 0; +} + +export namespace RenderedMessageOutputAreaCompat { + export interface IOptions { + chatPanel: ChatWidget; + } +} diff --git a/src/tools/commands.ts b/src/tools/commands.ts index eb2de654..7450fb9b 100644 --- a/src/tools/commands.ts +++ b/src/tools/commands.ts @@ -1,4 +1,5 @@ import { CommandRegistry } from '@lumino/commands'; +import { Widget } from '@lumino/widgets'; import { tool } from 'ai'; import { z } from 'zod'; import { ITool } from '../tokens'; @@ -115,9 +116,10 @@ export function createExecuteCommandTool( // Execute the command const result = await commands.execute(commandId, args); - // Handle Widget objects specially by extracting id and title + // Handle actual Lumino widgets specially by extracting id and title. + // Avoid collapsing plain command results that happen to contain an `id` field. let serializedResult; - if (result && typeof result === 'object' && result.id) { + if (result instanceof Widget) { serializedResult = { id: result.id, title: result.title?.label || result.title diff --git a/src/widgets/ai-settings.tsx b/src/widgets/ai-settings.tsx index 4ce7f6ab..5e4a64bc 100644 --- a/src/widgets/ai-settings.tsx +++ b/src/widgets/ai-settings.tsx @@ -31,7 +31,6 @@ import { InputLabel, List, ListItem, - ListItemSecondaryAction, ListItemText, Menu, MenuItem, @@ -1141,9 +1140,10 @@ const AISettingsComponent: React.FC = ({ {config.commandsRequiringApproval.map((command, index) => ( - - - + { const newCommands = [ @@ -1158,7 +1158,9 @@ const AISettingsComponent: React.FC = ({ > - + } + > + ))} @@ -1192,6 +1194,154 @@ const AISettingsComponent: React.FC = ({ )} /> + + + + + + {trans.__('Commands Auto-Rendering MIME Bundles')} + + + {trans.__( + 'Only these execute_command command IDs can auto-render MIME bundle outputs in chat' + )} + + + + {(config.commandsAutoRenderMimeBundles ?? []).map( + (command, index) => ( + { + const newCommands = [ + ...(config.commandsAutoRenderMimeBundles ?? + []) + ]; + newCommands.splice(index, 1); + handleConfigUpdate({ + commandsAutoRenderMimeBundles: newCommands + }); + }} + size="small" + > + + + } + > + + + ) + )} + + + { + if (e.key === 'Enter') { + const value = ( + e.target as HTMLInputElement + ).value.trim(); + const existingCommands = + config.commandsAutoRenderMimeBundles ?? []; + if (value && !existingCommands.includes(value)) { + const newCommands = [...existingCommands, value]; + handleConfigUpdate({ + commandsAutoRenderMimeBundles: newCommands + }); + (e.target as HTMLInputElement).value = ''; + } + } + }} + helperText={trans.__( + 'Press Enter to add a command. Default: jupyterlab-ai-commands:execute-in-kernel' + )} + /> + + + + + + + {trans.__('Trusted MIME Types for Auto-Render')} + + + {trans.__( + 'When auto-rendering command outputs, these MIME types are marked trusted in chat' + )} + + + + {(config.trustedMimeTypesForAutoRender ?? []).map( + (mimeType, index) => ( + { + const newMimeTypes = [ + ...(config.trustedMimeTypesForAutoRender ?? + []) + ]; + newMimeTypes.splice(index, 1); + handleConfigUpdate({ + trustedMimeTypesForAutoRender: newMimeTypes + }); + }} + size="small" + > + + + } + > + + + ) + )} + + + { + if (e.key === 'Enter') { + const value = ( + e.target as HTMLInputElement + ).value.trim(); + const existingMimeTypes = + config.trustedMimeTypesForAutoRender ?? []; + if (value && !existingMimeTypes.includes(value)) { + const newMimeTypes = [...existingMimeTypes, value]; + handleConfigUpdate({ + trustedMimeTypesForAutoRender: newMimeTypes + }); + (e.target as HTMLInputElement).value = ''; + } + } + }} + helperText={trans.__( + 'Press Enter to add a MIME type. Default: text/html' + )} + /> + @@ -1239,7 +1389,18 @@ const AISettingsComponent: React.FC = ({ ) : ( {config.mcpServers.map(server => ( - + handleMCPMenuClick(e, server.id)} + size="small" + > + + + } + > = ({ } /> - - handleMCPMenuClick(e, server.id)} - size="small" - > - - - ))} diff --git a/src/widgets/main-area-chat.ts b/src/widgets/main-area-chat.ts index 84a24ecf..ae13406a 100644 --- a/src/widgets/main-area-chat.ts +++ b/src/widgets/main-area-chat.ts @@ -8,6 +8,7 @@ import { ApprovalButtons } from '../approval-buttons'; import { AIChatModel } from '../chat-model'; import { TokenUsageWidget } from '../components/token-usage-display'; import { AISettingsModel } from '../models/settings-model'; +import { RenderedMessageOutputAreaCompat } from '../rendered-message-outputarea'; import { CommandIds } from '../tokens'; export namespace MainAreaChat { @@ -56,12 +57,18 @@ export class MainAreaChat extends MainAreaWidget { chatPanel: this.content, agentManager: this.model.agentManager }); + // Temporary compat: keep output-area CSS context for MIME renderers + // until jupyter-chat provides it natively. + this._outputAreaCompat = new RenderedMessageOutputAreaCompat({ + chatPanel: this.content + }); } dispose(): void { super.dispose(); // Dispose of the approval buttons widget when the chat is disposed. this._approvalButtons.dispose(); + this._outputAreaCompat.dispose(); } /** @@ -72,4 +79,5 @@ export class MainAreaChat extends MainAreaWidget { } private _approvalButtons: ApprovalButtons; + private _outputAreaCompat: RenderedMessageOutputAreaCompat; } diff --git a/ui-tests/tests/mime-bundles.spec.ts b/ui-tests/tests/mime-bundles.spec.ts new file mode 100644 index 00000000..44754ec7 --- /dev/null +++ b/ui-tests/tests/mime-bundles.spec.ts @@ -0,0 +1,174 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +import { + expect, + galata, + IJupyterLabPageFixture, + test +} from '@jupyterlab/galata'; +import { DEFAULT_GENERIC_PROVIDER_SETTINGS, openChatPanel } from './test-utils'; + +const EXPECT_TIMEOUT = 120000; +const TEST_MIME_BUNDLE_COMMAND_ID = 'jupyterlite-ai-tests:emit-mime-bundle'; +const BASE_SETTINGS = + DEFAULT_GENERIC_PROVIDER_SETTINGS['@jupyterlite/ai:settings-model']; +const PROVIDERS = BASE_SETTINGS.providers.map(provider => { + if (provider.id !== 'generic-functiongemma') { + return provider; + } + return { + ...provider, + parameters: { + ...(provider as { parameters?: Record }).parameters, + temperature: 0 + } + }; +}); + +test.use({ + mockSettings: { + ...galata.DEFAULT_SETTINGS, + '@jupyterlab/apputils-extension:notification': { + checkForUpdates: false, + fetchNews: 'false', + doNotDisturbMode: true + }, + '@jupyterlite/ai:settings-model': { + ...BASE_SETTINGS, + providers: PROVIDERS, + toolsEnabled: true, + defaultProvider: 'generic-functiongemma', + commandsAutoRenderMimeBundles: [TEST_MIME_BUNDLE_COMMAND_ID], + systemPrompt: + 'When asked to execute a command, call execute_command exactly once with this input shape: {"commandId":"jupyterlite-ai-tests:emit-mime-bundle"} and no args. Do not call any other tools and do not ask follow-up questions.' + } + } +}); + +async function registerTestMimeBundleCommand( + page: IJupyterLabPageFixture +): Promise { + await page.evaluate( + ({ commandId }) => { + const app = window.jupyterapp; + + if (app.commands.hasCommand(commandId)) { + return; + } + + app.commands.addCommand(commandId, { + label: 'Emit MIME bundle for UI tests', + execute: () => ({ + outputs: [ + { + output_type: 'display_data', + data: { + 'application/json': { + ok: true, + source: 'ui-test-command' + } + }, + metadata: {} + } + ] + }) + }); + }, + { commandId: TEST_MIME_BUNDLE_COMMAND_ID } + ); +} + +test.describe('#mimeBundles', () => { + test('should render MIME bundles from configured command outputs in chat', async ({ + page + }) => { + test.setTimeout(180 * 1000); + await registerTestMimeBundleCommand(page); + + const panel = await openChatPanel(page); + const input = panel + .locator('.jp-chat-input-container') + .getByRole('combobox'); + const sendButton = panel.locator( + '.jp-chat-input-container .jp-chat-send-button' + ); + + const prompt = `Call execute_command exactly once with commandId "${TEST_MIME_BUNDLE_COMMAND_ID}". Do not provide args.`; + + await input.pressSequentially(prompt); + await sendButton.click(); + + await expect( + panel.locator('.jp-chat-message-header:has-text("Jupyternaut")') + ).toHaveCount(1, { timeout: EXPECT_TIMEOUT }); + + const executeToolCall = panel + .locator('.jp-ai-tool-call') + .filter({ hasText: TEST_MIME_BUNDLE_COMMAND_ID }); + await expect(executeToolCall).toHaveCount(1, { timeout: EXPECT_TIMEOUT }); + await expect(executeToolCall).toContainText('execute_command', { + timeout: EXPECT_TIMEOUT + }); + await expect(executeToolCall).toContainText( + `"commandId": "${TEST_MIME_BUNDLE_COMMAND_ID}"`, + { + timeout: EXPECT_TIMEOUT + } + ); + await expect(executeToolCall).not.toContainText('"args": "', { + timeout: EXPECT_TIMEOUT + }); + + const renderedJson = panel.locator( + '.jp-chat-rendered-message .jp-RenderedJSON' + ); + await expect(renderedJson).toHaveCount(1, { timeout: EXPECT_TIMEOUT }); + }); + + test('should allow follow-up messages after MIME bundle auto-render', async ({ + page + }) => { + test.setTimeout(180 * 1000); + await registerTestMimeBundleCommand(page); + + const panel = await openChatPanel(page); + const input = panel + .locator('.jp-chat-input-container') + .getByRole('combobox'); + const sendButton = panel.locator( + '.jp-chat-input-container .jp-chat-send-button' + ); + const renderedJson = panel.locator( + '.jp-chat-rendered-message .jp-RenderedJSON' + ); + + const prompt = `Call execute_command exactly once with commandId "${TEST_MIME_BUNDLE_COMMAND_ID}". Do not provide args.`; + + await input.pressSequentially(prompt); + await expect(sendButton).toBeEnabled({ timeout: EXPECT_TIMEOUT }); + await sendButton.click(); + await expect(renderedJson).toHaveCount(1, { timeout: EXPECT_TIMEOUT }); + + const stopButton = panel.getByTitle('Stop streaming'); + await expect(stopButton).toHaveCount(0, { timeout: EXPECT_TIMEOUT }); + + await input.click(); + await input.pressSequentially(prompt); + await expect(sendButton).toBeEnabled({ timeout: EXPECT_TIMEOUT }); + await sendButton.click(); + + const executeToolCalls = panel + .locator('.jp-ai-tool-call') + .filter({ hasText: TEST_MIME_BUNDLE_COMMAND_ID }); + await expect(executeToolCalls).toHaveCount(2, { timeout: EXPECT_TIMEOUT }); + await expect(renderedJson).toHaveCount(2, { timeout: EXPECT_TIMEOUT }); + await expect( + panel.locator( + '.jp-chat-message-content:has-text("Error generating response:")' + ) + ).toHaveCount(0); + }); +});