feat(CodeBlock): support matplotlib in code execution (#8069)

* feat(CodeBlock): support matplotlib in code execution * refactor: update output style and docs * refactor: use ImageViewer * refactor: manage service config, increase timeout and retry count * refactor: improve worker message logging * chore: upgrade pyodide to 0.28.0 * docs: fix typos
2025-12-28 13:31:32 +08:00 · 2025-07-21 21:19:06 +08:00 · 2025-07-21 21:19:06 +08:00 · 929f7445ed
commit 929f7445ed
parent 63c3937050
7 changed files with 309 additions and 127 deletions
--- a/docs/technical/code-execution.md
+++ b/docs/technical/code-execution.md
@ -0,0 +1,127 @@
+# 代码执行功能
+
+本文档说明了代码块的 Python 代码执行功能。该实现利用 [Pyodide][pyodide-link] 在浏览器环境中直接运行 Python 代码，并将其置于 Web Worker 中，以避免阻塞主 UI 线程。
+
+整个实现分为三个主要部分：UI 层、服务层和 Worker 层。
+
+## 执行流程图
+
+```mermaid
+sequenceDiagram
+    participant 用户
+    participant CodeBlockView (UI)
+    participant PyodideService (服务)
+    participant PyodideWorker (Worker)
+
+    用户->>CodeBlockView (UI): 点击“运行”按钮
+    CodeBlockView (UI)->>PyodideService (服务): 调用 runScript(code)
+    PyodideService (服务)->>PyodideWorker (Worker): 发送 postMessage({ id, python: code })
+    PyodideWorker (Worker)->>PyodideWorker (Worker): 加载 Pyodide 和相关包
+    PyodideWorker (Worker)->>PyodideWorker (Worker): (按需)注入垫片并合并代码
+    PyodideWorker (Worker)->>PyodideWorker (Worker): 执行合并后的 Python 代码
+    PyodideWorker (Worker)-->>PyodideService (服务): 返回 postMessage({ id, output })
+    PyodideService (服务)-->>CodeBlockView (UI): 返回 { text, image } 对象
+    CodeBlockView (UI)->>用户: 在状态栏中显示文本和/或图像输出
+```
+
+## 1. UI 层
+
+面向用户的代码执行组件是 [CodeBlockView][codeblock-view-link]。
+
+### 关键机制：
+
+- **运行按钮**：当代码块语言为 `python` 且 `codeExecution.enabled` 设置为 true 时，`CodeToolbar` 中会条件性地渲染一个“运行”按钮。
+- **事件处理**：运行按钮的 `onClick` 事件会触发 `handleRunScript` 函数。
+- **服务调用**：`handleRunScript` 调用 `pyodideService.runScript(code)`，将代码块中的 Python 代码传递给服务。
+- **状态管理与输出显示**：使用 `executionResult` 来管理所有执行输出，只要有任何结果（文本或图像），[StatusBar][statusbar-link] 组件就会被渲染以统一显示。
+
+```typescript
+// src/renderer/src/components/CodeBlockView/view.tsx
+const [executionResult, setExecutionResult] = useState<{ text: string; image?: string } | null>(null)
+
+const handleRunScript = useCallback(() => {
+  setIsRunning(true)
+  setExecutionResult(null)
+
+  pyodideService
+    .runScript(children, {}, codeExecution.timeoutMinutes * 60000)
+    .then((result) => {
+      setExecutionResult(result)
+    })
+    .catch((error) => {
+      console.error('Unexpected error:', error)
+      setExecutionResult({
+        text: `Unexpected error: ${error.message || 'Unknown error'}`
+      })
+    })
+    .finally(() => {
+      setIsRunning(false)
+    })
+}, [children, codeExecution.timeoutMinutes]);
+
+// ... 在 JSX 中
+{isExecutable && executionResult && (
+  <StatusBar>
+    {executionResult.text}
+    {executionResult.image && (
+      <ImageOutput>
+        <img src={executionResult.image} alt="Matplotlib plot" />
+      </ImageOutput>
+    )}
+  </StatusBar>
+)}
+```
+
+## 2. 服务层
+
+服务层充当 UI 组件和运行 Pyodide 的 Web Worker 之间的桥梁。其逻辑封装在位于单例类 [PyodideService][pyodide-service-link]。
+
+### 主要职责：
+
+- **Worker 管理**：初始化、管理并与 Pyodide Web Worker 通信。
+- **请求处理**：使用 `resolvers` Map 管理并发请求，通过唯一 ID 匹配请求和响应。
+- **为 UI 提供 API**：向 UI 提供 `runScript(script, context, timeout)` 方法。此方法的返回值已修改为 `Promise<{ text: string; image?: string }>`，以支持包括图像在内的多种输出类型。
+- **输出处理**：从 Worker 接收包含文本、错误和可选图像数据的 `output` 对象。它将文本和错误格式化为对用户友好的单个字符串，然后连同图像数据一起包装成对象返回给 UI 层。
+- **IPC 端点**：该服务还提供了一个 `python-execution-request` IPC 端点，允许主进程请求执行 Python 代码，展示了其灵活的架构。
+
+## 3. Worker 层
+
+核心的 Python 执行发生在 [pyodide.worker.ts][pyodide-worker-link] 中定义的 Web Worker 内部。这确保了计算密集的 Python 代码不会冻结用户界面。
+
+### Worker 逻辑：
+
+- **Pyodide 加载**：Worker 从 CDN 加载 Pyodide 引擎，并设置处理器以捕获 Python 的 `stdout` 和 `stderr`。
+- **动态包安装**：使用 `pyodide.loadPackagesFromImports()` 自动分析并安装代码中导入的依赖包。
+- **按需执行垫片代码**：Worker 会检查传入的代码中是否包含 "matplotlib" 字符串。如果是，它会先执行一段 Python“垫片”代码确保图像输出到全局命名空间。
+- **结果序列化**：执行结果通过 `.toJs()` 等方法被递归转换为可序列化的标准 JavaScript 对象。
+- **返回结构化输出**：执行后，Worker 将一个包含 `id` 和 `output` 对象的-消息发回服务层。`output` 对象是一个结构化对象，包含 `result`、`text`、`error` 以及一个可选的 `image` 字段（用于 Base64 图像数据）。
+
+### 数据流
+
+最终的数据流如下：
+
+1.  **UI 层 ([CodeBlockView][codeblock-view-link])**: 用户点击“运行”按钮。
+2.  **服务层 ([PyodideService][pyodide-service-link])**:
+    - 接收到代码执行请求。
+    - 调用 Web Worker ([pyodide.worker.ts][pyodide-worker-link])，传递用户代码。
+3.  **Worker 层 ([pyodide.worker.ts][pyodide-worker-link])**:
+    - 加载 Pyodide 运行时。
+    - 动态安装代码中 `import` 语句声明的依赖包。
+    - **注入 Matplotlib 垫片**: 如果代码中包含 `matplotlib`，则在用户代码前拼接垫片代码，强制使用 `AGG` 后端。
+    - **执行代码并捕获输出**: 在代码执行后，检查 `matplotlib.pyplot` 的所有 figure，如果存在图像，则将其保存到内存中的 `BytesIO` 对象，并编码为 Base64 字符串。
+    - **结构化返回**: 将捕获的文本输出和 Base64 图像数据封装在一个 JSON 对象中 (`{ "text": "...", "image": "data:image/png;base64,..." }`) 返回给主线程。
+4.  **服务层 ([PyodideService][pyodide-service-link])**:
+    - 接收来自 Worker 的结构化数据。
+    - 将数据原样传递给 UI 层。
+5.  **UI 层 ([CodeBlockView][codeblock-view-link])**:
+    - 接收包含文本和图像数据的对象。
+    - 使用一个 `useState` 来管理执行结果 (`executionResult`)。
+    - 在界面上分别渲染文本输出和图像（如果存在）。
+
+<!-- Link Definitions -->
+
+[pyodide-link]: https://pyodide.org/
+[codeblock-view-link]: /src/renderer/src/components/CodeBlockView/view.tsx
+[pyodide-service-link]: /src/renderer/src/services/PyodideService.ts
+[pyodide-worker-link]: /src/renderer/src/workers/pyodide.worker.ts
+[statusbar-link]: /src/renderer/src/components/CodeBlockView/StatusBar.tsx
--- a/src/renderer/src/components/CodeBlockView/StatusBar.tsx
+++ b/src/renderer/src/components/CodeBlockView/StatusBar.tsx
@ -1,20 +1,21 @@
-import { FC, memo } from 'react'
+import { Flex } from 'antd'
+import { FC, memo, ReactNode } from 'react'
 import styled from 'styled-components'

 interface Props {
-  children: string
+  children: string | ReactNode
 }

 const StatusBar: FC<Props> = ({ children }) => {
  return <Container>{children}</Container>
 }

-const Container = styled.div`
-  margin: 10px;
+const Container = styled(Flex)`
+  background-color: var(--color-background-mute);
+  padding: 12px;
  display: flex;
-  flex-direction: row;
+  flex-direction: column;
  gap: 8px;
-  padding-bottom: 10px;
  overflow-y: auto;
  text-wrap: wrap;
 `
--- a/src/renderer/src/components/CodeBlockView/view.tsx
+++ b/src/renderer/src/components/CodeBlockView/view.tsx
@ -12,6 +12,7 @@ import React, { memo, useCallback, useEffect, useMemo, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'

+import ImageViewer from '../ImageViewer'
 import CodePreview from './CodePreview'
 import { SPECIAL_VIEW_COMPONENTS, SPECIAL_VIEWS } from './constants'
 import HtmlArtifactsCard from './HtmlArtifactsCard'
@ -48,7 +49,7 @@ export const CodeBlockView: React.FC<Props> = memo(({ children, language, onSave

  const [viewMode, setViewMode] = useState<ViewMode>('special')
  const [isRunning, setIsRunning] = useState(false)
-  const [output, setOutput] = useState('')
+  const [executionResult, setExecutionResult] = useState<{ text: string; image?: string } | null>(null)

  const [tools, setTools] = useState<CodeTool[]>([])
  const { registerTool, removeTool } = useCodeTool(setTools)
@ -87,16 +88,18 @@ export const CodeBlockView: React.FC<Props> = memo(({ children, language, onSave

  const handleRunScript = useCallback(() => {
    setIsRunning(true)
-    setOutput('')
+    setExecutionResult(null)

    pyodideService
      .runScript(children, {}, codeExecution.timeoutMinutes * 60000)
-      .then((formattedOutput) => {
-        setOutput(formattedOutput)
+      .then((result) => {
+        setExecutionResult(result)
      })
      .catch((error) => {
        logger.error('Unexpected error:', error)
-        setOutput(`Unexpected error: ${error.message || 'Unknown error'}`)
+        setExecutionResult({
+          text: `Unexpected error: ${error.message || 'Unknown error'}`
+        })
      })
      .finally(() => {
        setIsRunning(false)
@ -241,7 +244,14 @@ export const CodeBlockView: React.FC<Props> = memo(({ children, language, onSave
      {renderHeader}
      <CodeToolbar tools={tools} />
      {renderContent}
-      {isExecutable && output && <StatusBar>{output}</StatusBar>}
+      {isExecutable && executionResult && (
+        <StatusBar>
+          {executionResult.text}
+          {executionResult.image && (
+            <ImageViewer src={executionResult.image} alt="Matplotlib plot" style={{ cursor: 'pointer' }} />
+          )}
+        </StatusBar>
+      )}
    </CodeBlockWrapper>
  )
 })
--- a/src/renderer/src/pages/home/Markdown/ImagePreview.tsx
+++ b/src/renderer/src/pages/home/Markdown/ImagePreview.tsx
@ -1,64 +0,0 @@
-import {
-  DownloadOutlined,
-  RotateLeftOutlined,
-  RotateRightOutlined,
-  SwapOutlined,
-  UndoOutlined,
-  ZoomInOutlined,
-  ZoomOutOutlined
-} from '@ant-design/icons'
-import { download } from '@renderer/utils/download'
-import { Image as AntImage, ImageProps as AntImageProps, Space } from 'antd'
-import React from 'react'
-import styled from 'styled-components'
-
-interface ImagePreviewProps extends AntImageProps {
-  src: string
-}
-
-const ImagePreview: React.FC<ImagePreviewProps> = ({ src, ...props }) => {
-  return (
-    <AntImage
-      src={src}
-      {...props}
-      preview={{
-        mask: typeof props.preview === 'object' ? props.preview.mask : false,
-        toolbarRender: (
-          _,
-          {
-            transform: { scale },
-            actions: { onFlipY, onFlipX, onRotateLeft, onRotateRight, onZoomOut, onZoomIn, onReset }
-          }
-        ) => (
-          <ToobarWrapper size={12} className="toolbar-wrapper">
-            <SwapOutlined rotate={90} onClick={onFlipY} />
-            <SwapOutlined onClick={onFlipX} />
-            <RotateLeftOutlined onClick={onRotateLeft} />
-            <RotateRightOutlined onClick={onRotateRight} />
-            <ZoomOutOutlined disabled={scale === 1} onClick={onZoomOut} />
-            <ZoomInOutlined disabled={scale === 50} onClick={onZoomIn} />
-            <UndoOutlined onClick={onReset} />
-            <DownloadOutlined onClick={() => download(src)} />
-          </ToobarWrapper>
-        )
-      }}
-    />
-  )
-}
-
-const ToobarWrapper = styled(Space)`
-  padding: 0px 24px;
-  color: #fff;
-  font-size: 20px;
-  background-color: rgba(0, 0, 0, 0.1);
-  border-radius: 100px;
-  .anticon {
-    padding: 12px;
-    cursor: pointer;
-  }
-  .anticon:hover {
-    opacity: 0.3;
-  }
-`
-
-export default ImagePreview
--- a/src/renderer/src/services/PyodideService.ts
+++ b/src/renderer/src/services/PyodideService.ts
@ -3,11 +3,27 @@ import { uuid } from '@renderer/utils'

 const logger = loggerService.withContext('PyodideService')

+const SERVICE_CONFIG = {
+  WORKER: {
+    MAX_INIT_RETRY: 5, // 最大初始化重试次数
+    REQUEST_TIMEOUT: {
+      INIT: 30000, // 30 秒初始化超时
+      RUN: 60000 // 60 秒默认运行超时
+    }
+  }
+}
+
 // 定义结果类型接口
 export interface PyodideOutput {
  result: any
  text: string | null
  error: string | null
+  image?: string
+}
+
+export interface PyodideExecutionResult {
+  text: string
+  image?: string
 }

 /**
@ -19,7 +35,6 @@ class PyodideService {
  private worker: Worker | null = null
  private initPromise: Promise<void> | null = null
  private initRetryCount: number = 0
-  private static readonly MAX_INIT_RETRY = 2
  private resolvers: Map<string, { resolve: (value: any) => void; reject: (error: Error) => void }> = new Map()

  private constructor() {
@ -46,7 +61,7 @@ class PyodideService {
    if (this.worker) {
      return Promise.resolve()
    }
-    if (this.initRetryCount >= PyodideService.MAX_INIT_RETRY) {
+    if (this.initRetryCount >= SERVICE_CONFIG.WORKER.MAX_INIT_RETRY) {
      return Promise.reject(new Error('Pyodide worker initialization failed too many times'))
    }

@ -65,7 +80,7 @@ class PyodideService {
            this.initPromise = null
            this.initRetryCount++
            reject(new Error('Pyodide initialization timeout'))
-          }, 10000) // 10秒初始化超时
+          }, SERVICE_CONFIG.WORKER.REQUEST_TIMEOUT.INIT)

          // 设置初始化处理器
          const initHandler = (event: MessageEvent) => {
@ -75,7 +90,7 @@ class PyodideService {
              this.initRetryCount = 0
              this.initPromise = null
              resolve()
-            } else if (event.data?.type === 'error') {
+            } else if (event.data?.type === 'init-error') {
              clearTimeout(timeout)
              this.worker?.removeEventListener('message', initHandler)
              this.worker?.terminate()
@ -103,8 +118,16 @@ class PyodideService {
   * 处理来自 Worker 的消息
   */
  private handleMessage(event: MessageEvent): void {
+    const { type, error } = event.data
+
+    // 记录 Worker 错误消息
+    if (type === 'system-error') {
+      logger.error(error)
+      return
+    }
+
    // 忽略初始化消息，已由专门的处理器处理
-    if (event.data?.type === 'initialized' || event.data?.type === 'error') {
+    if (type === 'initialized' || type === 'init-error') {
      return
    }

@ -125,17 +148,23 @@ class PyodideService {
   * @param timeout 超时时间（毫秒）
   * @returns 格式化后的执行结果
   */
-  public async runScript(script: string, context: Record<string, any> = {}, timeout: number = 60000): Promise<string> {
+  public async runScript(
+    script: string,
+    context: Record<string, any> = {},
+    timeout: number = SERVICE_CONFIG.WORKER.REQUEST_TIMEOUT.RUN
+  ): Promise<PyodideExecutionResult> {
    // 确保Pyodide已初始化
    try {
      await this.initialize()
    } catch (error: unknown) {
      logger.error('Pyodide initialization failed, cannot execute Python code', error)
-      return `Initialization failed: ${error instanceof Error ? error.message : String(error)}`
+      const text = `Initialization failed: ${error instanceof Error ? error.message : String(error)}`
+      return { text }
    }

    if (!this.worker) {
-      return 'Internal error: Pyodide worker is not initialized'
+      const text = 'Internal error: Pyodide worker is not initialized'
+      return { text }
    }

    try {
@ -166,9 +195,10 @@ class PyodideService {
        })
      })

-      return this.formatOutput(output)
+      return { text: this.formatOutput(output), image: output.image }
    } catch (error: unknown) {
-      return `Internal error: ${error instanceof Error ? error.message : String(error)}`
+      const text = `Internal error: ${error instanceof Error ? error.message : String(error)}`
+      return { text }
    }
  }

@ -200,7 +230,7 @@ class PyodideService {
    // 如果有错误信息，附加显示
    if (output.error) {
      if (displayText) displayText += '\n\n'
-      displayText += `Error: ${output.error.trim()}`
+      displayText += output.error.trim()
    }

    // 如果没有任何输出，提供清晰提示
@ -250,13 +280,13 @@ if (typeof window !== 'undefined' && window.electron?.ipcRenderer) {

  window.electron.ipcRenderer.on('python-execution-request', async (_, request: PythonExecutionRequest) => {
    try {
-      const result = await pyodideService.runScript(request.script, request.context, request.timeout)
+      const { text } = await pyodideService.runScript(request.script, request.context, request.timeout)
      const response: PythonExecutionResponse = {
        id: request.id,
-        result
+        result: text
      }
      window.electron.ipcRenderer.send('python-execution-response', response)
-    } catch (error) {
+    } catch (error: unknown) {
      const response: PythonExecutionResponse = {
        id: request.id,
        error: error instanceof Error ? error.message : String(error)
--- a/src/renderer/src/services/ShikiStreamService.ts
+++ b/src/renderer/src/services/ShikiStreamService.ts
@ -13,6 +13,30 @@ import { ShikiStreamTokenizer, ShikiStreamTokenizerOptions } from './ShikiStream

 const logger = loggerService.withContext('ShikiStreamService')

+const SERVICE_CONFIG = {
+  // LRU 缓存配置
+  TOKENIZER_CACHE: {
+    MAX_SIZE: 100, // 最大缓存数量
+    TTL: 1000 * 60 * 30 // 30 分钟过期时间（毫秒）
+  },
+
+  // 降级策略配置
+  DEGRADATION_CACHE: {
+    MAX_SIZE: 500, // 最大记录数量
+    TTL: 1000 * 60 * 60 * 12 // 12 小时自动过期（毫秒）
+  },
+
+  // Worker 初始化配置
+  WORKER: {
+    MAX_INIT_RETRY: 2, // 最大初始化重试次数
+    REQUEST_TIMEOUT: {
+      INIT: 5000, // 初始化操作超时时间（毫秒）
+      HIGHLIGHT: 30000, // 高亮操作超时时间（毫秒）
+      DEFAULT: 10000 // 默认超时时间（毫秒）
+    }
+  }
+}
+
 export type ShikiPreProperties = {
  class: string
  style: string
@ -42,8 +66,8 @@ class ShikiStreamService {

  // 保存以 callerId-language-theme 为键的 tokenizer map
  private tokenizerCache = new LRUCache<string, ShikiStreamTokenizer>({
-    max: 100, // 最大缓存数量
-    ttl: 1000 * 60 * 30, // 30分钟过期时间
+    max: SERVICE_CONFIG.TOKENIZER_CACHE.MAX_SIZE,
+    ttl: SERVICE_CONFIG.TOKENIZER_CACHE.TTL,
    updateAgeOnGet: true,
    dispose: (value) => {
      if (value) value.clear()
@ -52,8 +76,8 @@ class ShikiStreamService {

  // 缓存每个 callerId 对应的已处理内容
  private codeCache = new LRUCache<string, string>({
-    max: 100, // 最大缓存数量
-    ttl: 1000 * 60 * 30, // 30分钟过期时间
+    max: SERVICE_CONFIG.TOKENIZER_CACHE.MAX_SIZE,
+    ttl: SERVICE_CONFIG.TOKENIZER_CACHE.TTL,
    updateAgeOnGet: true
  })

@ -61,7 +85,6 @@ class ShikiStreamService {
  private worker: Worker | null = null
  private workerInitPromise: Promise<void> | null = null
  private workerInitRetryCount: number = 0
-  private static readonly MAX_WORKER_INIT_RETRY = 2
  private pendingRequests = new Map<
    number,
    {
@ -73,8 +96,8 @@ class ShikiStreamService {

  // 降级策略相关变量，用于记录调用 worker 失败过的 callerId
  private workerDegradationCache = new LRUCache<string, boolean>({
-    max: 1000, // 最大记录数量
-    ttl: 1000 * 60 * 60 * 12 // 12小时自动过期
+    max: SERVICE_CONFIG.DEGRADATION_CACHE.MAX_SIZE,
+    ttl: SERVICE_CONFIG.DEGRADATION_CACHE.TTL
  })

  constructor() {
@ -103,7 +126,7 @@ class ShikiStreamService {
    if (this.workerInitPromise) return this.workerInitPromise
    if (this.worker) return

-    if (this.workerInitRetryCount >= ShikiStreamService.MAX_WORKER_INIT_RETRY) {
+    if (this.workerInitRetryCount >= SERVICE_CONFIG.WORKER.MAX_INIT_RETRY) {
      logger.debug('ShikiStream worker initialization failed too many times, stop trying')
      return
    }
@ -191,13 +214,13 @@ class ShikiStreamService {
      const getTimeoutForMessageType = (type: string): number => {
        switch (type) {
          case 'init':
-            return 5000 // 初始化操作 (5秒)
+            return SERVICE_CONFIG.WORKER.REQUEST_TIMEOUT.INIT
          case 'highlight':
-            return 30000 // 高亮操作 (30秒)
+            return SERVICE_CONFIG.WORKER.REQUEST_TIMEOUT.HIGHLIGHT
          case 'cleanup':
          case 'dispose':
          default:
-            return 10000 // 其他操作 (10秒)
+            return SERVICE_CONFIG.WORKER.REQUEST_TIMEOUT.DEFAULT
        }
      }

--- a/src/renderer/src/workers/pyodide.worker.ts
+++ b/src/renderer/src/workers/pyodide.worker.ts
@ -1,16 +1,63 @@
 /// <reference lib="webworker" />

-import { loggerService } from '@logger'
-
-const logger = loggerService.initWindowSource('Worker').withContext('Pyodide')
+interface WorkerResponse {
+  type: 'initialized' | 'init-error' | 'system-error'
+  id?: string
+  output?: PyodideOutput
+  error?: string
+}

 // 定义输出结构类型
 interface PyodideOutput {
  result: any
  text: string | null
  error: string | null
+  image?: string
 }

+const PYODIDE_INDEX_URL = 'https://cdn.jsdelivr.net/pyodide/v0.28.0/full/'
+const PYODIDE_MODULE_URL = PYODIDE_INDEX_URL + 'pyodide.mjs'
+
+// 垫片代码，用于在 Worker 中捕获 Matplotlib 绘图
+const MATPLOTLIB_SHIM_CODE = `
+def __cherry_studio_matplotlib_setup():
+    import os
+    # 在导入 pyplot 前设置后端
+    os.environ["MPLBACKEND"] = "AGG"
+    import io
+    import base64
+    import matplotlib.pyplot as plt
+
+    # 保存原始的 show 函数
+    _original_show = plt.show
+
+    # 定义并替换为新的 show 函数
+    def _new_show(*args, **kwargs):
+        global pyodide_matplotlib_image
+        fig = plt.gcf()
+
+        if not fig.canvas.get_renderer()._renderer:
+            return
+
+        buf = io.BytesIO()
+        fig.savefig(buf, format='png')
+        buf.seek(0)
+
+        img_str = base64.b64encode(buf.read()).decode('utf-8')
+
+        # 通过全局变量传递数据
+        pyodide_matplotlib_image = f"data:image/png;base64,{img_str}"
+
+        plt.clf()
+        plt.close(fig)
+
+    # 替换全局的 show 函数
+    plt.show = _new_show
+
+__cherry_studio_matplotlib_setup()
+del __cherry_studio_matplotlib_setup
+`
+
 // 声明全局变量用于输出
 let output: PyodideOutput = {
  result: null,
@ -28,12 +75,11 @@ const pyodidePromise = (async () => {

  try {
    // 动态加载 Pyodide 脚本
-    // @ts-ignore - 忽略动态导入错误
-    const pyodideModule = await import('https://cdn.jsdelivr.net/pyodide/v0.27.5/full/pyodide.mjs')
+    const pyodideModule = await import(/* @vite-ignore */ PYODIDE_MODULE_URL)

    // 加载 Pyodide 并捕获标准输出/错误
    return await pyodideModule.loadPyodide({
-      indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.27.5/full/',
+      indexURL: PYODIDE_INDEX_URL,
      stdout: (text: string) => {
        if (output.text) {
          output.text += `${text}\n`
@ -51,13 +97,12 @@ const pyodidePromise = (async () => {
    })
  } catch (error: unknown) {
    const errorMessage = error instanceof Error ? error.message : String(error)
-    logger.error('Failed to load Pyodide:', errorMessage)

    // 通知主线程初始化错误
    self.postMessage({
-      type: 'error',
+      type: 'init-error',
      error: errorMessage
-    })
+    } as WorkerResponse)

    throw error
  }
@ -81,7 +126,6 @@ function processResult(result: any): any {
    return result
  } catch (error: unknown) {
    const errorMessage = error instanceof Error ? error.message : String(error)
-    logger.error('Result processing error:', errorMessage)
    return { __error__: 'Result processing failed', details: errorMessage }
  }
 }
@ -89,17 +133,19 @@ function processResult(result: any): any {
 // 通知主线程已加载
 pyodidePromise
  .then(() => {
-    self.postMessage({ type: 'initialized' })
+    self.postMessage({ type: 'initialized' } as WorkerResponse)
  })
  .catch((error: unknown) => {
    const errorMessage = error instanceof Error ? error.message : String(error)
-    logger.error('Failed to load Pyodide:', errorMessage)
-    self.postMessage({ type: 'error', error: errorMessage })
+    self.postMessage({
+      type: 'init-error',
+      error: errorMessage
+    } as WorkerResponse)
  })

 // 处理消息
 self.onmessage = async (event) => {
-  const { id, python, context } = event.data
+  const { id, python } = event.data

  // 重置输出变量
  output = {
@ -111,12 +157,6 @@ self.onmessage = async (event) => {
  try {
    const pyodide = await pyodidePromise

-    // 将上下文变量设置为全局作用域变量
-    const globalContext: Record<string, any> = {}
-    for (const key of Object.keys(context || {})) {
-      globalContext[key] = context[key]
-    }
-
    // 载入需要的包
    try {
      await pyodide.loadPackagesFromImports(python)
@ -125,14 +165,23 @@ self.onmessage = async (event) => {
      throw new Error(`Failed to load required packages: ${errorMessage}`)
    }

-    // 创建 Python 上下文
-    const globals = pyodide.globals.get('dict')(Object.entries(context || {}))
-
    // 执行代码
    try {
-      output.result = await pyodide.runPythonAsync(python, { globals })
+      // 注入 Matplotlib 垫片代码
+      if (python.includes('matplotlib')) {
+        await pyodide.runPythonAsync(MATPLOTLIB_SHIM_CODE)
+      }
+
+      output.result = await pyodide.runPythonAsync(python)
      // 处理结果，确保安全序列化
      output.result = processResult(output.result)
+
+      // 检查是否有 Matplotlib 图像输出
+      const image = pyodide.globals.get('pyodide_matplotlib_image')
+      if (image) {
+        output.image = image
+        pyodide.globals.delete('pyodide_matplotlib_image')
+      }
    } catch (error: unknown) {
      const errorMessage = error instanceof Error ? error.message : String(error)
      // 不设置 output.result，但设置错误信息
@ -145,15 +194,21 @@ self.onmessage = async (event) => {
  } catch (error: unknown) {
    // 处理所有其他错误
    const errorMessage = error instanceof Error ? error.message : String(error)
-    logger.error('Python processing error:', errorMessage)

    if (output.error) {
      output.error += `\nSystem error:\n${errorMessage}`
    } else {
      output.error = `System error:\n${errorMessage}`
    }
+
+    // 发送错误信息
+    self.postMessage({
+      type: 'system-error',
+      id,
+      error: errorMessage
+    } as WorkerResponse)
  } finally {
    // 统一发送处理后的输出对象
-    self.postMessage({ id, output })
+    self.postMessage({ id, output } as WorkerResponse)
  }
 }