notebook-ai-extension/src/context/cell_context.ts

/**
 * Extract a CellContext snapshot from a CodeCell + its parent NotebookPanel.
 *
 * Pure functions, no DOM, no signals — testable with plain object mocks.
 */
import type { CodeCell, ICodeCellModel } from '@jupyterlab/cells';
import type { NotebookPanel } from '@jupyterlab/notebook';

import type { CellContext, ErrorOutput } from '../types';

/**
 * Find the first error output in a code cell's outputs.
 * Returns null if the cell has no error output or is not a code cell.
 */
export function collectErrorOutputs(model: ICodeCellModel): ErrorOutput | null {
  const outputs = model.outputs;
  for (let i = 0; i < outputs.length; i++) {
    const out = outputs.get(i);
    if (out && out.type === 'error') {
      const e = out as {
        ename?: unknown;
        evalue?: unknown;
        traceback?: unknown;
      };
      return {
        ename: String(e.ename ?? ''),
        evalue: String(e.evalue ?? ''),
        traceback: Array.isArray(e.traceback) ? (e.traceback as string[]) : [],
      };
    }
  }
  return null;
}

/**
 * Build a CellContext from a CodeCell and its parent NotebookPanel.
 * Returns null if essential data is missing.
 *
 * @param cell The code cell to inspect.
 * @param panel The notebook panel that contains the cell.
 */
export function extractCellContext(
  cell: CodeCell,
  panel: NotebookPanel
): CellContext | null {
  if (!cell || !panel) {
    return null;
  }

  const cells = panel.content.widgets;
  const cellIndex = cells ? cells.indexOf(cell) : -1;
  if (cellIndex < 0) {
    return null;
  }

  const totalCells = cells ? cells.length : 0;
  const previousCell = cellIndex > 0 ? cells[cellIndex - 1] : null;
  const previousCode = previousCell
    ? previousCell.model.sharedModel.getSource()
    : null;

  // Language: use 'python' as default for code cells.
  // Future improvement: read from kernel info_reply.
  const language = 'python';

  let error: ErrorOutput | null = null;
  if (cell.model.type === 'code') {
    error = collectErrorOutputs(cell.model as ICodeCellModel);
  }

  return {
    notebookPath: panel.context.path,
    cellId: cell.model.id,
    language,
    cellIndex,
    totalCells,
    source: cell.model.sharedModel.getSource(),
    previousCode,
    error,
  };
}