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

/**
 * Global floating AI button + panel, mounted on document.body. Replaces
 * the per-cell toolbar AI button (formerly OpenCodeCellActions). Click
 * the round bottom-right button to expand the OpenCodeInlinePrompt; click
 * again to collapse it.
 *
 * The panel operates on the JupyterLab shell's currently-active
 * NotebookPanel — the (notebookPath, activeCell) pair is resolved at
 * panel-open time from `app.shell.currentWidget`, and re-resolved when
 * `app.shell.currentChanged` fires while the panel is open.
 *
 * All session/SSE/history state that used to live on
 * OpenCodeCellActions now lives here. The prompt is created lazily on
 * the first panel open and disposed on close (recreated on every
 * re-open so the user always sees a fresh history fetch).
 */
import { JupyterFrontEnd } from '@jupyterlab/application';
import { Notification } from '@jupyterlab/apputils';
import type { CodeCell } from '@jupyterlab/cells';
import { ServerConnection } from '@jupyterlab/services';
import { Widget } from '@lumino/widgets';

import {
  callOpenCodeBindExistingSession,
  callOpenCodeCreateNotebookSession,
  callOpenCodeEdit,
  callOpenCodeListAllSessions,
  callOpenCodeReplyPermission,
  callOpenCodeReplyQuestion,
  callOpenCodeSessionMessages,
  subscribeOpenCodeEvents,
  type OpenCodeEventSubscription
} from '../api/opencode_client';
import type { OpenCodeProvidersResponse } from '../types';

import { OpenCodeInlinePrompt } from './opencode_inline_prompt';

export interface IOpenCodeFloatingPanelOptions {
  app: JupyterFrontEnd;
  serverSettings: ServerConnection.ISettings;
  providers: OpenCodeProvidersResponse | null;
}

type Status = 'idle' | 'loading';

export class OpenCodeFloatingPanel extends Widget {
  private _app: JupyterFrontEnd;
  private _serverSettings: ServerConnection.ISettings | null;
  private _providers: OpenCodeProvidersResponse | null;
  private _button: HTMLButtonElement;
  private _panelEl: HTMLDivElement;
  private _prompt: OpenCodeInlinePrompt | null = null;
  private _sseSub: OpenCodeEventSubscription | null = null;
  private _notebookPath: string | null = null;
  private _status: Status = 'idle';

  constructor(options: IOpenCodeFloatingPanelOptions) {
    super();
    this._app = options.app;
    this._serverSettings = options.serverSettings;
    this._providers = options.providers;
    this.addClass('opencode-floating-root');
    this.node.style.position = 'fixed';
    this.node.style.bottom = '24px';
    this.node.style.right = '24px';
    this.node.style.zIndex = '1000';

    this._button = document.createElement('button');
    this._button.className = 'opencode-floating-button';
    this._button.textContent = '🪄';
    this._button.title = 'OpenCode AI';
    this._button.addEventListener('click', () => this.toggle());
    this.node.appendChild(this._button);

    this._panelEl = document.createElement('div');
    this._panelEl.className = 'opencode-floating-panel';
    this._panelEl.hidden = true;
    this.node.appendChild(this._panelEl);

    // JupyterFrontEnd.shell is typed as `JupyterShell | null` in some
    // JupyterLab versions; cast through any to satisfy strict null checks
    // while keeping the public API simple.
    (this._app.shell as any).currentChanged.connect(this._onShellChanged, this);
  }

  dispose(): void {
    if (this.isDisposed) {
      return;
    }
    (this._app.shell as any).currentChanged.disconnect(
      this._onShellChanged,
      this
    );
    this._closeSse();
    this._hidePrompt();
    super.dispose();
  }

  /** Update the providers payload after the initial async fetch. */
  setProviders(p: OpenCodeProvidersResponse | null): void {
    this._providers = p;
  }

  /** Toggle the panel open/closed. */
  toggle(): void {
    if (this._panelEl.hidden) {
      this._showPrompt();
    } else {
      this._hidePrompt();
    }
  }

  private _onShellChanged = (): void => {
    // The active notebook changed. If the panel is open, refresh the
    // underlying prompt so the (provider, model) selection, history,
    // and active-cell context all reflect the new notebook.
    //
    // Implemented as an arrow-function class field (not a method) so
    // `this` is bound regardless of how the signal's slot invokes us.
    if (!this._panelEl.hidden) {
      this._hidePrompt();
      this._showPrompt();
    }
  };

  private _resolveNotebookContext(): {
    notebookPath: string | null;
    getActiveCell: () => CodeCell | null;
  } {
    const widget = this._app.shell.currentWidget as any;
    const notebookPath =
      widget && widget.context && typeof widget.context.path === 'string'
        ? (widget.context.path as string)
        : null;
    const getActiveCell = (): CodeCell | null => {
      const cur = this._app.shell.currentWidget as any;
      if (!cur || !cur.content) {
        return null;
      }
      const active = cur.content.activeCell;
      return (active as CodeCell) || null;
    };
    return { notebookPath, getActiveCell };
  }

  private _showPrompt(): void {
    if (this._prompt) {
      this._panelEl.hidden = false;
      return;
    }
    const { notebookPath, getActiveCell } = this._resolveNotebookContext();
    this._notebookPath = notebookPath;
    const prompt = new OpenCodeInlinePrompt({
      disabled: !this._notebookPath || this._status === 'loading',
      providers: this._providers,
      notebookPath: this._notebookPath ?? undefined,
      getActiveCell,
      onSubmit: (text: string, providerId?: string, modelId?: string) => {
        void this._onSubmit(text, providerId, modelId);
      },
      onCancel: () => {
        this._hidePrompt();
      },
      onPermissionReply: async (permId, response) => {
        if (!this._serverSettings) {
          return false;
        }
        const sessionId = (this._prompt as any)._sessionId as string | null;
        if (!sessionId) {
          return false;
        }
        try {
          const r = await callOpenCodeReplyPermission(
            sessionId,
            permId,
            response,
            this._serverSettings
          );
          return r.ok;
        } catch (e) {
          console.warn('opencode_bridge: permission reply failed', e);
          return false;
        }
      },
      onQuestionReply: async (qId, answer) => {
        if (!this._serverSettings) {
          return false;
        }
        const sessionId = (this._prompt as any)._sessionId as string | null;
        if (!sessionId) {
          return false;
        }
        try {
          const r = await callOpenCodeReplyQuestion(
            sessionId,
            qId,
            answer,
            this._serverSettings
          );
          return r.ok;
        } catch (e) {
          console.warn('opencode_bridge: question reply failed', e);
          return false;
        }
      },
      onStreamEnd: () => {
        // Centralized: the prompt detected "idle" and signals us to
        // close the SSE stream + restore the input.
        this._closeSse();
      },
      onListSessions: async () => {
        if (!this._serverSettings) {
          return [];
        }
        try {
          const r = await callOpenCodeListAllSessions(this._serverSettings);
          return r.sessions || [];
        } catch (e) {
          console.warn('opencode_bridge: list all sessions failed', e);
          return [];
        }
      },
      onCreateSession: async () => {
        if (!this._serverSettings || !this._notebookPath) {
          throw new Error('notebook not ready');
        }
        this._closeSse();
        const r = await callOpenCodeCreateNotebookSession(
          this._notebookPath,
          undefined,
          this._serverSettings
        );
        if (!r.ok || !r.sessionId) {
          throw new Error(r.error || 'create_session failed');
        }
        return r.sessionId;
      },
      onSwitchSession: async (sessionId: string) => {
        if (!this._serverSettings || !this._notebookPath) {
          return false;
        }
        this._closeSse();
        // PUT /sessions/notebook binds AND sets active in one call.
        const r = await callOpenCodeBindExistingSession(
          this._notebookPath,
          sessionId,
          this._serverSettings
        );
        return !!r.ok;
      },
      onReloadHistory: async () => {
        if (!this._notebookPath) {
          return;
        }
        await this._refreshHistory(this._notebookPath);
      }
    });
    // Use Widget.attach (not raw appendChild) so Lumino sends the
    // before-attach / after-attach messages — the prompt relies on
    // these for isAttached/isVisible tracking. The defensive DOM
    // cleanup in _hidePrompt handles the case where dispose() leaves
    // a residual node behind (the panel host is an HTMLElement, not a
    // Widget, so the Lumino parent stays null and dispose() can't
    // always clean up the DOM on its own).
    Widget.attach(prompt, this._panelEl);
    this._prompt = prompt;

    // Load the current session's static history before streaming starts.
    if (this._notebookPath) {
      void this._refreshHistory(this._notebookPath);
    }
    prompt.initialize();
    this._panelEl.hidden = false;
  }

  private async _refreshHistory(notebookPath: string): Promise<void> {
    if (!this._serverSettings || !this._prompt) {
      return;
    }
    try {
      const resp = await callOpenCodeSessionMessages(
        notebookPath,
        this._serverSettings
      );
      this._prompt.setMessages(resp.messages);
    } catch (e) {
      console.warn('opencode_bridge: failed to fetch session messages', e);
    }
  }

  private _hidePrompt(): void {
    this._closeSse();
    if (this._prompt) {
      const node = this._prompt.node;
      this._prompt.dispose();
      // Defensive DOM cleanup: the prompt is attached via raw DOM
      // (appendChild/Widget.attach with an HTMLElement host), so its
      // Lumino parent is null. In some scenarios — for example when
      // the dispose path runs synchronously while another show+hide
      // cycle is already in flight — the dispose() call does not
      // remove the node from this._panelEl, leaving a residual node
      // that the next _showPrompt appends to. Without this defensive
      // removeChild, the user would see the panel content multiply
      // with every click.
      if (node.parentNode === this._panelEl) {
        this._panelEl.removeChild(node);
      }
      this._prompt = null;
    }
    this._panelEl.hidden = true;
  }

  private async _onSubmit(
    text: string,
    providerId?: string,
    modelId?: string
  ): Promise<void> {
    if (!this._notebookPath) {
      Notification.info('请先打开一个 notebook');
      return;
    }
    if (!this._serverSettings) {
      Notification.error('OpenCode 运行时未初始化');
      return;
    }

    const request = {
      prompt: text,
      context: { notebookPath: this._notebookPath },
      providerId,
      modelId
    };

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

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

    if (!resp.ok) {
      this._status = 'idle';
      this._prompt?.setDisabled(false);
      Notification.error(`OpenCode 错误: ${resp.error}`);
      return;
    }

    this._prompt?.setSessionId(resp.sessionId);
    this._prompt?.clearInput();
    this._sseSub = subscribeOpenCodeEvents(
      resp.sessionId,
      this._serverSettings,
      {
        onEvent: ev => {
          this._prompt?.applyEvent(ev);
        },
        onError: err => {
          console.warn('opencode_bridge: SSE error', err);
        }
      }
    );
  }

  private _closeSse(): void {
    if (this._sseSub) {
      this._sseSub.close();
      this._sseSub = null;
    }
    this._status = 'idle';
    this._prompt?.setDisabled(false);
  }
}