frontend/juypter/ai-cell-generator: tricky off-by-one error needs to be fixed ...

Author: haraldschilly

src/packages/frontend/jupyter/insert-cell/ai-cell-generator.tsx

@@ -196,20 +196,26 @@ export function AIGenerateCodeCell({
   }, [preview, open]);
 
   function getContextContents(): CellContextContent {
-    const prevCount = -contextRange[0]; // contextRange[0] is negative, so -(-2) = 2
-    const nextCount = contextRange[1]; // contextRange[1] is positive for cells after
+    const [rangeStart, rangeEnd] = contextRange;
 
-    if (prevCount === 0 && nextCount === 0) return {};
+    // For insertion mode:
+    // - rangeStart to rangeEnd defines which cells relative to insertion point to include
+    // - Current cell id represents position 0 (insertion point)
+    // - [-2, 0] means include cells from id-1 to id (2 cells before insertion)
+    // - [0, 2] means include cells from id+1 to id+2 (2 cells after insertion)
+
+    const aboveCount = rangeStart < 0 ? Math.abs(rangeStart) : 0;
+    const belowCount = rangeEnd > 0 ? rangeEnd : 0;
 
     return getNonemptyCellContents({
       actions: frameActions.current,
       id,
       direction: "around",
-      cellCount: "all", // Use "all" for around direction
+      cellCount: "all",
       cellTypes,
       lang,
-      aboveCount: prevCount,
-      belowCount: nextCount,
+      aboveCount,
+      belowCount,
     });
   }
 

src/packages/frontend/jupyter/llm/cell-context-selector.tsx

@@ -42,25 +42,49 @@ export function LLMCellContextSelector({
       return { minValue: 0, maxValue: 0, marks: { 0: "0" } };
     }
 
-    // Count cells above
-    let cellsAbove = 0;
-    let delta = -1;
-    while (jupyterActionsStore.get_cell_id(delta, currentCellId)) {
-      cellsAbove++;
-      delta--;
-    }
-
-    // Count cells below
-    let cellsBelow = 0;
-    delta = 1;
-    while (jupyterActionsStore.get_cell_id(delta, currentCellId)) {
-      cellsBelow++;
-      delta++;
+    let minVal: number, maxVal: number;
+
+    if (mode === "insert-position") {
+      // For insert position, we need to count cells differently
+      // Count all cells before the insertion point (include current cell and all above it)
+      let cellsBefore = 0;
+      let delta = 0; // Start from current cell (which will be "before" after insertion)
+      while (jupyterActionsStore.get_cell_id(delta, currentCellId)) {
+        cellsBefore++;
+        delta--;
+      }
+
+      // Count cells after
+      let cellsAfter = 0;
+      delta = 1;
+      while (jupyterActionsStore.get_cell_id(delta, currentCellId)) {
+        cellsAfter++;
+        delta++;
+      }
+
+      minVal = -cellsBefore;
+      maxVal = cellsAfter;
+    } else {
+      // For current-cell mode, count cells above and below as before
+      let cellsAbove = 0;
+      let delta = -1;
+      while (jupyterActionsStore.get_cell_id(delta, currentCellId)) {
+        cellsAbove++;
+        delta--;
+      }
+
+      // Count cells below
+      let cellsBelow = 0;
+      delta = 1;
+      while (jupyterActionsStore.get_cell_id(delta, currentCellId)) {
+        cellsBelow++;
+        delta++;
+      }
+
+      minVal = -cellsAbove;
+      maxVal = cellsBelow;
     }
 
-    const minVal = -cellsAbove;
-    const maxVal = cellsBelow;
-
     // Create marks dynamically
     const marks: SliderSingleProps["marks"] = {
       0: mode === "current-cell" ? "0" : "insert",

src/packages/frontend/jupyter/util/cell-content.ts

@@ -70,6 +70,7 @@ export function getNonemptyCellContents({
         lang,
         direction: "below",
         count: belowCount,
+        includeCurrentCell: false,
       });
       if (belowContent) result.after = belowContent;
     }
@@ -82,10 +83,10 @@ export function getNonemptyCellContents({
     typeof cellCount === "number"
       ? cellCount
       : cellCount === "all above" ||
-          cellCount === "all below" ||
-          cellCount === "all"
-        ? 100
-        : 0;
+        cellCount === "all below" ||
+        cellCount === "all"
+      ? 100
+      : 0;
 
   const content = getDirectionalContent({
     actions,