notebook-ai-extension/design.md

📑 1. 系统总体架构

采用 前端 UI 注入 + 本地 Server Extension 透传 的轻量双层架构：

┌────────────────────────────────────────────────────────────────────────┐
│                   JupyterLab / Notebook v7 (前端 TS)                   │
│                                                                        │
│   ┌────────────────────────────────────────────────────────────────┐   │
│   │                         Cell Component                         │   │
│   │   [CodeEditor] ─── (挂载) ───► [Cell Action Toolbar]           │   │
│   │                                       │ (点击 🪄 智能编辑)          │
│   │                                       ▼                            │
│   │                             [Inline Prompt Box]                    │   │
│   └───────────────────────────────────┬────────────────────────────┘   │
│                                       │                                │
│                            Context Provider 收集                      │
│                            (Target Cell, Error, Neighbors)             │
│                                       │                                │
└───────────────────────────────────────┼────────────────────────────────┘
                                        │ REST / SSE Stream
                                        ▼
┌────────────────────────────────────────────────────────────────────────┐
│                 Jupyter Server Extension (后端 Python)                 │
│                                                                        │
│   ┌────────────────────────────────────────────────────────────────┐   │
│   │                     OpenCode Bridge Handler                    │   │
│   │   - 接收 Context & Prompt                                      │   │
│   │   - 拼装标准 OpenCode API Payload                              │   │
│   │   - SSE 流式透传代理 (Stream Proxy)                            │   │
│   └───────────────────────────────────┬────────────────────────────┘   │
└───────────────────────────────────────┼────────────────────────────────┘
                                        │ Local Socket / HTTP
                                        ▼
┌────────────────────────────────────────────────────────────────────────┐
│                           OpenCode Serve 进程                          │
└────────────────────────────────────────────────────────────────────────┘

---

🎨 2. 交互与UI流程设计

1. 原生 Cell Toolbar 单按钮（2026-07-23 v2 修正）：不再自绘悬浮 toolbar，也合并了 v1 的三个 mode（优化/排错/编辑）为单个 🪄 AI 智能编辑 图标按钮。通过 IToolbarWidgetRegistry.addFactory('Cell', 'opencode-cell-actions', ...) 注入 JupyterLab 原生 cell toolbar，显示在 active cell 右上角（move up/down 按钮旁，rank: 100 排在 delete 右侧）；仅 CodeCell 显示，随 active cell 切换移动。前端不做任何语义处理 —— 不再发 mode 字段，后端用统一的系统提示词 + 用户的自然语言指令 + 上下文（含 traceback）由 LLM 自己判断是优化/排错/编辑。

2. Inline Prompt 输入框（v2 实现）：点击 🪄 AI 智能编辑 按钮后，OpenCodeInlinePrompt widget 追加到当前 CodeCell 的 DOM 末尾，呈现内嵌输入面板：多行 textarea（placeholder "向 AI 描述你想要的修改…"，无默认值）+ 🚀 发送 + ✕ 取消 两个按钮。提交时前端仅收集 cell source / outputs / error 与用户在 textarea 里写的指令，POST /opencode-bridge/edit，body 为 {prompt, context, providerId?, modelId?}（无 mode）。成功 → 前端调用 cell.model.sharedModel.setSource(resp.finalSource) 就地替换 cell 源码 + Notification.info；失败 → Notification.error，输入框保留（可改文案重发）。Session 由服务端 SessionManager 保证 1 notebook 1 sid（前端不感知）。

3. 流式替换/Diff 预览：

• 点击发送后，支持打字机实时替换或在 Cell 下方出现 [ Accept ] / [ Reject ] 临时 Diff 比对面板。

---

🔌 3. 接口契约设计 (API Contract)

接口：POST /opencode-bridge/edit（v2 — 无 mode）

Request Payload（前端 OpenCodeCellActions 在用户提交 inline 输入框时发出）：

{
  "prompt": "把这段代码改为用 plotly 绘制交互式折线图",
  "context": {
    "notebookPath": "analysis/demo.ipynb",
    "cellId": "cell-3",
    "language": "python",
    "cellIndex": 3,
    "totalCells": 5,
    "source": "import matplotlib.pyplot as plt\nplt.plot([1, 2, 3])",
    "previousCode": "import pandas as pd",
    "error": null
  },
  "providerId": "anthropic",
  "modelId": "claude-sonnet-4-20250514"
}

v2 变化：body 不再含 mode / action 字段。前端只把 cell 的 input/output/error 收集到 context，加用户的自然语言 prompt 一起发出；后端用统一的 UNIFIED_SYSTEM_PROMPT + LLM 自己判断优化/排错/编辑。providerId / modelId 来自 JupyterLab settings 透传。

Response（当前 v2 实现：单次 JSON；SSE 流式属于 Slice 4 后续工作）：

{
  "ok": true,
  "finalSource": "import plotly.express as px\nfig = px.line(x=[1, 2, 3], y=[1, 2, 3])\nfig.show()",
  "sessionId": "ses-abc123",
  "notebookPath": "analysis/demo.ipynb"
}

成功时前端用 finalSource 就地替换 cell 源码（cell.model.sharedModel.setSource）。Session 由服务端 SessionManager 按 notebookPath 复用（1 notebook 1 sid），404 / session not found 时服务端自动 invalidate 并在下次请求重建。

---

🛠️ 4. 核心前端实现逻辑 (TypeScript)

以下是前端实现“Toolbar 注入”和“上下文抓取”的伪代码骨架，可直接作为开发参考：

src/context_provider.ts (上下文感知模块)

import { Notebook, NotebookTracker } from '@jupyterlab/notebook';
import { Cell } from '@jupyterlab/cells';

export interface CellContext {
  notebookPath: string;
  targetCode: string;
  errorOutput?: string;
  surroundingCells: string[];
}

export function extractCellContext(tracker: NotebookTracker, targetCell: Cell): CellContext | null {
  const currentWidget = tracker.currentWidget;
  if (!currentWidget) return null;

  const notebook: Notebook = currentWidget.content;
  const cells = notebook.widgets;
  const targetIndex = cells.findIndex(c => c === targetCell);

  // 1. 抓取当前 Cell 的错误输出
  let errorMsg = '';
  if (targetCell.model.type === 'code') {
    const outputs = (targetCell.model as any).outputs;
    for (let i = 0; i < outputs.length; i++) {
      const out = outputs.get(i);
      if (out.type === 'error') {
        errorMsg += `${out.ename}: ${out.evalue}\n${(out.traceback || []).join('\n')}`;
      }
    }
  }

  // 2. 抓取前一个 Cell 的代码（上下文）
  const prevCode = targetIndex > 0 ? cells[targetIndex - 1].model.sharedModel.getSource() : '';

  return {
    notebookPath: currentWidget.context.path,
    targetCode: targetCell.model.sharedModel.getSource(),
    errorOutput: errorMsg,
    surroundingCells: [prevCode]
  };
}

src/components/cell_toolbar.ts (工具栏与内嵌框注入)

import { Cell } from '@jupyterlab/cells';

export function injectOpenCodeActionBox(cell: Cell, onSendPrompt: (prompt: string) => void) {
  // 避免重复注入
  if (cell.node.querySelector('.opencode-toolbar')) return;

  // 1. 创建 Toolbar 容器
  const toolbar = document.createElement('div');
  toolbar.className = 'opencode-toolbar';
  toolbar.innerHTML = `
    <button class="opencode-btn edit-btn">🪄 智能编辑</button>
  `;

  // 2. 绑定展开事件
  toolbar.querySelector('.edit-btn')!.addEventListener('click', () => {
    let promptBox = cell.node.querySelector('.opencode-inline-box') as HTMLDivElement;
    if (promptBox) {
      promptBox.style.display = promptBox.style.display === 'none' ? 'block' : 'none';
      return;
    }

    // 动态创建 Inline Prompt 输入框
    promptBox = document.createElement('div');
    promptBox.className = 'opencode-inline-box';
    promptBox.innerHTML = `
      <textarea placeholder="输入提示词（例如：重构这段代码或添加注释...）"></textarea>
      <div class="actions">
        <button class="submit-btn">🚀 发送给 OpenCode</button>
      </div>
    `;

    promptBox.querySelector('.submit-btn')!.addEventListener('click', () => {
      const input = promptBox.querySelector('textarea')!.value;
      if (input.trim()) onSendPrompt(input);
    });

    cell.node.appendChild(promptBox);
  });

  cell.node.insertBefore(toolbar, cell.node.firstChild);
}

---

🔗 5. Session 管理策略（v0.2.1 修正）

2026-07-22 修正：原 v0.2 设计 EditHandler 采用"每请求一个 session + finally 清理"——这是错的。OpenCode 的 session 是有状态的对话上下文，每个 notebook 必须复用同一个 session，否则用户先优化 cell1、再编辑 cell2 时，AI 看不到前一次的结果。

5.1 核心规则

维度	决策
映射关系	notebookPath（1 个）→ OpenCode sessionID（1 个）
生命周期	首次 EditHandler 调用时 lazy create；显式 DELETE /session 才释放
并发	同 notebook 的并发请求用 asyncio.Lock 串行化 session 创建；消息发送并发无锁（OpenCode 端自己排队）
失效恢复	当 send_message_sync 抛 OpenCodeError 且消息含 404/410 → 自动 invalidate 缓存，下次请求重新 create
清理	不做自动释放。前端 notebook 关闭时不调 release（避免两个 tab 共享 session 时误杀）。用户可手动点"重置对话"触发 release
跨重启	Jupyter server 重启 → in-memory map 丢失 → OpenCode 端孤立 session 接受泄漏（v0.4 加启动清理）

5.2 SessionManager 接口

class SessionManager:
    def __init__(self, client_factory: Callable[[], OpenCodeClient]): ...

    async def get_or_create(self, notebook_path: str) -> str:
        """Return existing sessionID, or create one. Async-safe via per-path lock."""

    async def release(self, notebook_path: str) -> bool:
        """Delete session and remove from map. Returns True if session existed."""

    def has_session(self, notebook_path: str) -> bool: ...

    def list_sessions(self) -> list[dict]:
        """For debug: [{notebookPath, sessionId}, ...]"""

关键不变量：

• get_or_create 永远不返回 None 或抛 "not found"
• 重复调用 get_or_create(same_path) → 同一 sessionID
• release(path) 后再 get_or_create(path) → 新 sessionID

5.3 HTTP 端点

Method	Path	作用	备注
GET	/opencode-bridge/sessions	列出所有 active session	debug 用
DELETE	/opencode-bridge/session?notebook=<path>	显式释放指定 notebook 的 session	URL-encode path

为什么用 query param 而不是 path param：

• notebook path 可能含 /、空格、中文
• url_path_join + regex 捕获在含 / 的路径上行为不可控
• query string 是 RFC 3986 标准做法，tornado 自动 decode

5.4 EditHandler 改造

改造前（错的）：

create_session → send_message_sync → delete_session   # 每个请求都这样

改造后（对的）：

session_id = session_manager.get_or_create(notebook_path)   # 首次创建，复用后续
send_message_sync(session_id, ...)                          # 不管 session
# 没有 delete — session 留着给下次

5.5 错误恢复

send_message_sync 失败时：

1. 检查 OpenCodeError 消息是否含 "404" 或 "session not found"
2. 是 → session_manager.invalidate(notebook_path)，下次 get_or_create 会重新创建
3. 否 → 透传错误给前端

v0.2.1 不实现自动重试，只 invalidate。前端看到 502 后可重发。

5.6 未解决问题（v0.4+）

1. TTL reaper：notebook 关闭后 session 永远不释放 → 改用 mtime 追踪 + 30 分钟 idle 自动 release
2. 多 server 横向扩展：in-memory state 不能 scale → 改成 Redis 或 sticky session
3. 启动清理：Jupyter 重启后调 GET /session 列出 OpenCode 端的孤儿，逐个 delete
4. 并发 LLM 调用：当前两次 send_message_sync 到同一 session 是串行（OpenCode 端排队）— 可加 asyncio.Semaphore(per_session=1) 显式控制

---

🎛 6. Provider/Model 选择（v0.2.2 新增）

2026-07-22：用户要求能在 settings 里选 OpenCode 的 provider 和 model，可选值从 GET /config/providers 动态获取。

6.1 Schema 扩展

schema/plugin.json 加 2 个 string 字段：

"opencodeProvider": {
  "type": "string",
  "title": "OpenCode Provider",
  "description": "Provider id (如 'anthropic' / 'openai')。可选项由启动时从 /opencode-bridge/providers 拉取并打到 console 列出。留空 = 用 OpenCode 默认。",
  "default": ""
},
"opencodeModel": {
  "type": "string",
  "title": "OpenCode Model",
  "description": "Model id (如 'claude-sonnet-4-20250514')。可选项见 console 日志或 /opencode-bridge/providers。留空 = 用 provider 默认。",
  "default": ""
}

6.2 选值如何给用户

JupyterLab settings 是静态 JSON schema，没有原生 dynamic enum。v0.2.2 用最简方案：

• 启动时 fetch /opencode-bridge/providers → 把 provider/model 列表 console.log 出来
• 用户在 Settings Editor 里手填
• 后续 v0.4 考虑用 QuickPick 命令做交互式选择

6.3 数据流

Settings (JupyterLab)
  │
  ├── opencodeProvider = "anthropic"
  └── opencodeModel    = "claude-sonnet-4-20250514"
        │
        ▼  setOpenCodeRuntime()
  Frontend (opencode_cell_footer.ts)
        │
        ▼  build request body
  POST /opencode-bridge/edit
  { mode, prompt, context, providerId, modelId }
        │
        ▼  send_message_sync
  OpenCode Serve 接受 model = { providerID, modelID }

6.4 不变量

• 留空字段（""）→ 客户端不发 providerId/modelId → server 也不加 model 字段到 OpenCode 请求 → OpenCode 用默认
• 用户在 settings 里设的值不会自动回写到 OpenCode——只是 request 级别的覆盖
• 改 settings 立即生效（下次 edit 请求就用新值）；不需要重启 Jupyter