notebook-ai-extension/src/components/opencode_cell_actions.ts

/**
 * Per-cell AI action button rendered inside JupyterLab's native Cell toolbar
 * (top-right of the active cell). v2: a single icon button. Clicking it
 * toggles an OpenCodeInlinePrompt panel attached to the cell's DOM.
 *
 * Registered in index.ts via IToolbarWidgetRegistry.addFactory('Cell', ...).
 *
 * v3-final: the frontend does NO semantic processing. It gathers the cell's
 * source / outputs / error, attaches the user's freeform instruction, and
 * POSTs to the server. The chosen provider/model comes from the inline
 * picker (no settings-based default).
 */
import type { CodeCell } from '@jupyterlab/cells';
import { Notification } from '@jupyterlab/apputils';
import type { NotebookPanel } from '@jupyterlab/notebook';
import { ServerConnection } from '@jupyterlab/services';
import { Widget } from '@lumino/widgets';

import { callOpenCodeEdit, callOpenCodeSessionMessages } from '../api/opencode_client';
import { extractCellContext } from '../context/cell_context';
import type {
  CellContext,
  OpenCodeProvidersResponse,
  OpenCodeRequest,
  OpenCodeResponse
} from '../types';

import { OpenCodeInlinePrompt } from './opencode_inline_prompt';

// Module-level runtime injection (set by index.ts after activation).
let _serverSettings: ServerConnection.ISettings | null = null;
let _providers: OpenCodeProvidersResponse | null = null;

export function setOpenCodeServerSettings(
  serverSettings: ServerConnection.ISettings
): void {
  _serverSettings = serverSettings;
}

export function setOpenCodeProviders(
  p: OpenCodeProvidersResponse | null
): void {
  _providers = p;
}

type Status = 'idle' | 'loading';

export class OpenCodeCellActions extends Widget {
  private _cell: CodeCell;
  private _context: CellContext | null = null;
  private _status: Status = 'idle';
  private _prompt: OpenCodeInlinePrompt | null = null;

  constructor(cell: CodeCell) {
    super();
    this._cell = cell;
    this.addClass('opencode-cell-actions');
    this._cell.model.contentChanged.connect(this._onModelChange, this);
    this._cell.model.stateChanged.connect(this._onModelChange, this);
    this._render();
    this._onModelChange();
  }

  dispose(): void {
    if (this.isDisposed) {
      return;
    }
    this._cell.model.contentChanged.disconnect(this._onModelChange, this);
    this._cell.model.stateChanged.disconnect(this._onModelChange, this);
    this._hidePrompt();
    super.dispose();
  }

  private _onModelChange(): void {
    this._context = extractCellContextFromCell(this._cell);
    if (this._prompt) {
      this._prompt.setDisabled(!this._context || this._status === 'loading');
    } else {
      this._render();
    }
  }

  private _render(): void {
    const node = this.node;
    node.textContent = '';

    const baseDisabled = !this._context || this._status === 'loading';

    const btn = document.createElement('button');
    btn.className = 'opencode-btn opencode-btn-ai';
    btn.textContent = 'AI';
    btn.disabled = baseDisabled;
    btn.title = baseDisabled
      ? 'OpenCode: 等待 cell 上下文…'
      : '让 AI 修改这个 cell 的代码';
    btn.addEventListener('click', () => {
      this._togglePrompt();
    });
    node.appendChild(btn);
  }

  private _togglePrompt(): void {
    if (this._prompt) {
      this._hidePrompt();
    } else {
      this._showPrompt();
    }
  }

  private _showPrompt(): void {
    if (this._prompt) {
      return;
    }
    const prompt = new OpenCodeInlinePrompt(this._cell, {
      disabled: !this._context || this._status === 'loading',
      providers: _providers,
      onSubmit: (text: string, providerId?: string, modelId?: string) => {
        void this._onSubmit(text, providerId, modelId);
      },
      onCancel: () => {
        this._hidePrompt();
      }
    });
    Widget.attach(prompt, this._cell.node);
    this._prompt = prompt;

    // Fetch the current session's message history and render it into
    // the prompt's scrollable history area. Empty list (no session yet)
    // is a normal no-op render.
    const notebookPath = this._context?.notebookPath;
    if (notebookPath) {
      void this._refreshHistory(notebookPath);
    }
  }

  private async _refreshHistory(notebookPath: string): Promise<void> {
    if (!_serverSettings || !this._prompt) {
      return;
    }
    try {
      const resp = await callOpenCodeSessionMessages(
        notebookPath,
        _serverSettings
      );
      this._prompt.setMessages(resp.messages);
    } catch (e) {
      // Non-fatal: the history is a convenience. The user can still
      // send new prompts; just the scrollback won't update.
      console.warn('opencode_bridge: failed to fetch session messages', e);
    }
  }

  private _hidePrompt(): void {
    if (this._prompt) {
      this._prompt.dispose();
      this._prompt = null;
    }
  }

  private async _onSubmit(
    text: string,
    providerId?: string,
    modelId?: string
  ): Promise<void> {
    if (!this._context) {
      return;
    }
    if (!_serverSettings) {
      Notification.error('OpenCode 运行时未初始化');
      return;
    }

    const request: OpenCodeRequest = {
      prompt: text,
      context: this._context,
      providerId,
      modelId
    };

    this._status = 'loading';
    if (this._prompt) {
      this._prompt.setDisabled(true);
    }

    try {
      const resp = await callOpenCodeEdit(request, _serverSettings);
      this._handleResponse(resp);
    } catch (e) {
      this._status = 'idle';
      if (this._prompt) {
        this._prompt.setDisabled(false);
      }
      Notification.error(`OpenCode 失败: ${(e as Error).message}`);
    }
  }

  private _handleResponse(resp: OpenCodeResponse): void {
    this._status = 'idle';
    if (resp.ok) {
      // Refetch the (now-updated) session history and render the new
      // assistant message as the last item in the scrollable history.
      // The cell source is NOT replaced.
      const notebookPath = this._context?.notebookPath;
      if (notebookPath) {
        void this._refreshHistory(notebookPath);
      }
      // Clear the input so the user can type a follow-up message.
      this._prompt?.clearInput();
      Notification.info(
        `OpenCode 完成 (${resp.markdown.length} chars). Session: ${resp.sessionId.slice(0, 8)}…`
      );
    } else {
      if (this._prompt) {
        this._prompt.setDisabled(false);
      }
      Notification.error(`OpenCode 错误: ${resp.error}`);
    }
  }
}

/**
 * Resolve a CodeCell's parent NotebookPanel by walking the parent chain,
 * then build the CellContext.
 */
function extractCellContextFromCell(cell: CodeCell): CellContext | null {
  let node: Widget | null = cell.parent;
  let notebookPanel: NotebookPanel | null = null;
  while (node) {
    const candidate = node as any;
    if (
      candidate.context &&
      candidate.content &&
      Array.isArray(candidate.content.widgets)
    ) {
      notebookPanel = candidate;
      break;
    }
    node = node.parent;
  }
  if (!notebookPanel) {
    return null;
  }
  return extractCellContext(cell, notebookPanel);
}