Refactor inline comment architecture across plan/files/diff and fix diff overlay rendering (#461)

packages/ui/src/components/comments/CodeMirrorCommentWidgets.tsx

@@ -0,0 +1,106 @@
+import React from 'react';
+import { InlineCommentCard } from './InlineCommentCard';
+import { InlineCommentInput } from './InlineCommentInput';
+import type { InlineCommentDraft } from '@/stores/useInlineCommentDraftStore';
+import type { BlockWidgetDef } from '@/components/ui/CodeMirrorEditor';
+
+type LineRange = {
+  start: number;
+  end: number;
+  side?: 'additions' | 'deletions';
+};
+
+interface CodeMirrorCommentWidgetsOptions {
+  drafts: InlineCommentDraft[];
+  editingDraftId: string | null;
+  commentText: string;
+  selection: LineRange | null;
+  isDragging: boolean;
+  fileLabel: string;
+  newWidgetId: string;
+  mapDraftToRange: (draft: InlineCommentDraft) => LineRange;
+  onSave: (text: string, range?: LineRange) => void;
+  onCancel: () => void;
+  onEdit: (draft: InlineCommentDraft) => void;
+  onDelete: (draft: InlineCommentDraft) => void;
+}
+
+export function buildCodeMirrorCommentWidgets(options: CodeMirrorCommentWidgetsOptions): BlockWidgetDef[] {
+  const {
+    drafts,
+    editingDraftId,
+    commentText,
+    selection,
+    isDragging,
+    fileLabel,
+    newWidgetId,
+    mapDraftToRange,
+    onSave,
+    onCancel,
+    onEdit,
+    onDelete,
+  } = options;
+
+  const widgets: BlockWidgetDef[] = [];
+
+  for (const draft of drafts) {
+    const draftRange = mapDraftToRange(draft);
+    if (draft.id === editingDraftId) {
+      widgets.push({
+        afterLine: draftRange.end,
+        id: `edit-${draft.id}`,
+        content: (
+          <InlineCommentInput
+            key={`edit-${draft.id}`}
+            initialText={commentText}
+            fileLabel={fileLabel}
+            lineRange={draftRange}
+            isEditing={true}
+            onSave={onSave}
+            onCancel={onCancel}
+          />
+        ),
+      });
+      continue;
+    }
+
+    widgets.push({
+      afterLine: draftRange.end,
+      id: `card-${draft.id}`,
+      content: (
+        <InlineCommentCard
+          key={`card-${draft.id}`}
+          draft={draft}
+          onEdit={() => onEdit(draft)}
+          onDelete={() => onDelete(draft)}
+        />
+      ),
+    });
+  }
+
+  if (selection && !editingDraftId && !isDragging) {
+    const normalizedSelection = {
+      ...selection,
+      start: Math.min(selection.start, selection.end),
+      end: Math.max(selection.start, selection.end),
+    };
+
+    widgets.push({
+      afterLine: normalizedSelection.end,
+      id: newWidgetId,
+      content: (
+        <InlineCommentInput
+          key={newWidgetId}
+          initialText={commentText}
+          fileLabel={fileLabel}
+          lineRange={normalizedSelection}
+          isEditing={false}
+          onSave={onSave}
+          onCancel={onCancel}
+        />
+      ),
+    });
+  }
+
+  return widgets;
+}

packages/ui/src/components/comments/InlineCommentCard.tsx

@@ -17,6 +17,7 @@ interface InlineCommentCardProps {
   onEdit: () => void;
   onDelete: () => void;
   className?: string;
+  maxWidth?: number;
 }
 
 export function InlineCommentCard({
@@ -24,6 +25,7 @@ export function InlineCommentCard({
   onEdit,
   onDelete,
   className,
+  maxWidth,
 }: InlineCommentCardProps) {
   const themeContext = useOptionalThemeSystem();
   const currentTheme = themeContext?.currentTheme;
@@ -42,6 +44,7 @@ export function InlineCommentCard({
       style={{
         backgroundColor: currentTheme?.colors?.surface?.elevated,
         borderColor: currentTheme?.colors?.interactive?.border,
+        maxWidth: maxWidth ? `${Math.max(200, Math.floor(maxWidth))}px` : undefined,
       }}
       data-comment-card="true"
     >

packages/ui/src/components/comments/InlineCommentInput.tsx

@@ -13,6 +13,7 @@ export interface InlineCommentInputProps {
   lineRange?: { start: number; end: number; side?: 'additions' | 'deletions' };
   isEditing?: boolean;
   className?: string;
+  maxWidth?: number;
 }
 
 export function InlineCommentInput({
@@ -23,6 +24,7 @@ export function InlineCommentInput({
   lineRange,
   isEditing = false,
   className,
+  maxWidth,
 }: InlineCommentInputProps) {
   const themeContext = useOptionalThemeSystem();
   const currentTheme = themeContext?.currentTheme;
@@ -118,6 +120,7 @@ export function InlineCommentInput({
       style={{
         backgroundColor: currentTheme?.colors?.surface?.elevated,
         borderColor: currentTheme?.colors?.interactive?.border,
+        maxWidth: maxWidth ? `${Math.max(200, Math.floor(maxWidth))}px` : undefined,
       }}
       data-comment-input="true"
       onPointerDown={(e) => e.stopPropagation()}

packages/ui/src/components/comments/PierreDiffCommentOverlays.tsx

@@ -0,0 +1,230 @@
+import React from 'react';
+import { createPortal } from 'react-dom';
+import type { SelectedLineRange } from '@pierre/diffs';
+import type { InlineCommentDraft } from '@/stores/useInlineCommentDraftStore';
+import { InlineCommentCard } from './InlineCommentCard';
+import { InlineCommentInput } from './InlineCommentInput';
+import { toPierreAnnotationId } from './PierreDiffCommentUtils';
+
+interface PierreDiffCommentOverlaysProps {
+  diffRootRef: React.RefObject<HTMLDivElement | null>;
+  drafts: InlineCommentDraft[];
+  selection: SelectedLineRange | null;
+  editingDraftId: string | null;
+  commentText: string;
+  fileLabel: string;
+  onSave: (text: string, range?: SelectedLineRange) => void;
+  onCancel: () => void;
+  onEdit: (draft: InlineCommentDraft) => void;
+  onDelete: (draft: InlineCommentDraft) => void;
+}
+
+function parseCssWidth(value: string): number | null {
+  const trimmed = value.trim();
+  if (!trimmed) return null;
+
+  if (/^-?\d+(\.\d+)?$/.test(trimmed)) {
+    const parsed = Number.parseFloat(trimmed);
+    return Number.isFinite(parsed) ? parsed : null;
+  }
+
+  if (trimmed.endsWith('px')) {
+    const parsed = Number.parseFloat(trimmed.slice(0, -2));
+    return Number.isFinite(parsed) ? parsed : null;
+  }
+
+  return null;
+}
+
+function clampMaxWidth(value: number | null | undefined): number | undefined {
+  if (!value || value <= 0) return undefined;
+  return Math.max(200, Math.floor(value));
+}
+
+export function PierreDiffCommentOverlays(props: PierreDiffCommentOverlaysProps) {
+  const {
+    diffRootRef,
+    drafts,
+    selection,
+    editingDraftId,
+    commentText,
+    fileLabel,
+    onSave,
+    onCancel,
+    onEdit,
+    onDelete,
+  } = props;
+
+  const [retryTick, setRetryTick] = React.useState(0);
+  const [fallbackMaxWidth, setFallbackMaxWidth] = React.useState<number | null>(null);
+
+  const selectionAnnotationId = React.useMemo(() => {
+    if (!selection || editingDraftId) return null;
+    return toPierreAnnotationId({ type: 'new', selection });
+  }, [editingDraftId, selection]);
+
+  const expectedTargetIds = React.useMemo(() => {
+    const ids = drafts.map((draft) => toPierreAnnotationId({ type: draft.id === editingDraftId ? 'edit' : 'saved', draft }));
+    if (selectionAnnotationId) {
+      ids.push(selectionAnnotationId);
+    }
+    return ids;
+  }, [drafts, editingDraftId, selectionAnnotationId]);
+
+  const resolveTarget = React.useCallback((annotationId: string): HTMLElement | null => {
+    const wrapper = diffRootRef.current;
+    if (!wrapper) return null;
+
+    const host = wrapper.querySelector('diffs-container');
+    if (!(host instanceof HTMLElement)) return null;
+
+    const lightDomTarget = host.querySelector(`[data-annotation-id="${annotationId}"]`);
+    if (lightDomTarget instanceof HTMLElement) {
+      return lightDomTarget;
+    }
+
+    const shadowRoot = host.shadowRoot;
+    if (!shadowRoot) return null;
+
+    return shadowRoot.querySelector(`[data-annotation-id="${annotationId}"]`) as HTMLElement | null;
+  }, [diffRootRef]);
+
+  React.useEffect(() => {
+    if (expectedTargetIds.length === 0) return;
+
+    let cancelled = false;
+    let attempts = 0;
+    const maxAttempts = 12;
+
+    const checkTargets = () => {
+      if (cancelled) return;
+      const allResolved = expectedTargetIds.every((id) => Boolean(resolveTarget(id)));
+      if (allResolved || attempts >= maxAttempts) {
+        return;
+      }
+      attempts += 1;
+      requestAnimationFrame(() => {
+        if (cancelled) return;
+        setRetryTick((tick) => tick + 1);
+        checkTargets();
+      });
+    };
+
+    checkTargets();
+    return () => {
+      cancelled = true;
+    };
+  }, [expectedTargetIds, resolveTarget]);
+
+  React.useEffect(() => {
+    const root = diffRootRef.current;
+    if (!root) return;
+
+    const computeMaxWidth = () => {
+      const styles = getComputedStyle(root);
+      const cssWidth = parseCssWidth(styles.getPropertyValue('--oc-context-panel-width'));
+      const rootRect = root.getBoundingClientRect();
+      const measured = cssWidth ?? rootRect.width;
+      setFallbackMaxWidth(measured > 0 ? measured : null);
+    };
+
+    computeMaxWidth();
+
+    const observer = new ResizeObserver(() => {
+      computeMaxWidth();
+    });
+    observer.observe(root);
+
+    window.addEventListener('resize', computeMaxWidth);
+    return () => {
+      observer.disconnect();
+      window.removeEventListener('resize', computeMaxWidth);
+    };
+  }, [diffRootRef]);
+
+  const resolveTargetMaxWidth = React.useCallback((target: HTMLElement): number | undefined => {
+    const root = diffRootRef.current;
+    const rootRect = root?.getBoundingClientRect();
+
+    const annotationContent = target.closest('[data-annotation-content]');
+    const contentRect = annotationContent instanceof HTMLElement
+      ? annotationContent.getBoundingClientRect()
+      : target.getBoundingClientRect();
+
+    const candidates = [contentRect.width];
+    if (rootRect) {
+      candidates.push(rootRect.right - contentRect.left);
+    }
+
+    const positiveCandidates = candidates.filter((value) => Number.isFinite(value) && value > 0);
+    if (positiveCandidates.length > 0) {
+      return clampMaxWidth(Math.min(...positiveCandidates));
+    }
+
+    return clampMaxWidth(fallbackMaxWidth);
+  }, [diffRootRef, fallbackMaxWidth]);
+
+  void retryTick;
+
+  return (
+    <>
+      {drafts.map((draft) => {
+        const id = toPierreAnnotationId({ type: draft.id === editingDraftId ? 'edit' : 'saved', draft });
+        const target = resolveTarget(id);
+        if (!target) return null;
+        const targetMaxWidth = resolveTargetMaxWidth(target);
+
+        if (draft.id === editingDraftId) {
+          return createPortal(
+            <InlineCommentInput
+              initialText={commentText}
+              fileLabel={fileLabel}
+              lineRange={{
+                start: draft.startLine,
+                end: draft.endLine,
+                side: draft.side === 'original' ? 'deletions' : 'additions',
+              }}
+              isEditing={true}
+              onSave={onSave}
+              onCancel={onCancel}
+              maxWidth={targetMaxWidth}
+            />,
+            target,
+            `draft-edit-${draft.id}`
+          );
+        }
+
+        return createPortal(
+          <InlineCommentCard
+            draft={draft}
+            onEdit={() => onEdit(draft)}
+            onDelete={() => onDelete(draft)}
+            maxWidth={targetMaxWidth}
+          />,
+          target,
+          `draft-card-${draft.id}`
+        );
+      })}
+
+      {selection && !editingDraftId && selectionAnnotationId && (() => {
+        const target = resolveTarget(selectionAnnotationId);
+        if (!target) return null;
+        const targetMaxWidth = resolveTargetMaxWidth(target);
+
+        return createPortal(
+          <InlineCommentInput
+            initialText={commentText}
+            fileLabel={fileLabel}
+            lineRange={selection}
+            isEditing={false}
+            onSave={onSave}
+            onCancel={onCancel}
+            maxWidth={targetMaxWidth}
+          />,
+          target,
+          selectionAnnotationId
+        );
+      })()}
+    </>
+  );
+}

packages/ui/src/components/comments/PierreDiffCommentUtils.ts

@@ -0,0 +1,52 @@
+import type { AnnotationSide, DiffLineAnnotation, SelectedLineRange } from '@pierre/diffs';
+import type { InlineCommentDraft } from '@/stores/useInlineCommentDraftStore';
+
+export type PierreAnnotationData =
+  | { type: 'saved' | 'edit'; draft: InlineCommentDraft }
+  | { type: 'new'; selection: SelectedLineRange };
+
+export const toPierreAnnotationId = (meta: PierreAnnotationData): string => {
+  if (meta.type === 'new') {
+    const start = Math.min(meta.selection.start, meta.selection.end);
+    const end = Math.max(meta.selection.start, meta.selection.end);
+    const side = meta.selection.side ?? 'additions';
+    return `new-comment-${side}-${start}-${end}`;
+  }
+
+  return `draft-${meta.draft.id}`;
+};
+
+interface BuildPierreLineAnnotationsOptions {
+  drafts: InlineCommentDraft[];
+  editingDraftId: string | null;
+  selection: SelectedLineRange | null;
+}
+
+export const buildPierreLineAnnotations = (
+  options: BuildPierreLineAnnotationsOptions
+): DiffLineAnnotation<PierreAnnotationData>[] => {
+  const { drafts, editingDraftId, selection } = options;
+  const annotations: DiffLineAnnotation<PierreAnnotationData>[] = [];
+
+  for (const draft of drafts) {
+    const side: AnnotationSide = draft.side === 'original' ? 'deletions' : 'additions';
+    annotations.push({
+      lineNumber: draft.endLine,
+      side,
+      metadata: {
+        type: draft.id === editingDraftId ? 'edit' : 'saved',
+        draft,
+      },
+    });
+  }
+
+  if (selection && !editingDraftId) {
+    annotations.push({
+      lineNumber: Math.max(selection.start, selection.end),
+      side: selection.side ?? 'additions',
+      metadata: { type: 'new', selection },
+    });
+  }
+
+  return annotations;
+};

packages/ui/src/components/comments/index.ts

@@ -1,2 +1,6 @@
 export * from './InlineCommentCard';
 export * from './InlineCommentInput';
+export * from './useInlineCommentController';
+export * from './CodeMirrorCommentWidgets';
+export * from './PierreDiffCommentUtils';
+export * from './PierreDiffCommentOverlays';

packages/ui/src/components/comments/useInlineCommentController.ts

@@ -0,0 +1,154 @@
+import React from 'react';
+import { toast } from '@/components/ui';
+import { useInlineCommentDraftStore, type InlineCommentDraft, type InlineCommentSource } from '@/stores/useInlineCommentDraftStore';
+import { useSessionStore } from '@/stores/useSessionStore';
+
+type LineRangeBase = {
+  start: number;
+  end: number;
+};
+
+type StoreRange = {
+  startLine: number;
+  endLine: number;
+  side?: 'original' | 'modified';
+};
+
+interface UseInlineCommentControllerOptions<TRange extends LineRangeBase> {
+  source: InlineCommentSource;
+  fileLabel: string | null;
+  language: string;
+  getCodeForRange: (range: TRange) => string;
+  toStoreRange: (range: TRange) => StoreRange;
+  fromDraftRange: (draft: InlineCommentDraft) => TRange;
+}
+
+const normalizeStoreRange = (range: StoreRange): StoreRange => {
+  const startLine = Math.min(range.startLine, range.endLine);
+  const endLine = Math.max(range.startLine, range.endLine);
+  return {
+    ...range,
+    startLine,
+    endLine,
+  };
+};
+
+export const normalizeLineRange = <TRange extends LineRangeBase>(range: TRange): TRange => {
+  const start = Math.min(range.start, range.end);
+  const end = Math.max(range.start, range.end);
+  return {
+    ...range,
+    start,
+    end,
+  };
+};
+
+export function useInlineCommentController<TRange extends LineRangeBase>(
+  options: UseInlineCommentControllerOptions<TRange>
+) {
+  const { source, fileLabel, language, getCodeForRange, toStoreRange, fromDraftRange } = options;
+
+  const currentSessionId = useSessionStore((state) => state.currentSessionId);
+  const newSessionDraftOpen = useSessionStore((state) => state.newSessionDraft?.open);
+
+  const addDraft = useInlineCommentDraftStore((state) => state.addDraft);
+  const updateDraft = useInlineCommentDraftStore((state) => state.updateDraft);
+  const removeDraft = useInlineCommentDraftStore((state) => state.removeDraft);
+  const allDrafts = useInlineCommentDraftStore((state) => state.drafts);
+
+  const [selection, setSelection] = React.useState<TRange | null>(null);
+  const [commentText, setCommentText] = React.useState('');
+  const [editingDraftId, setEditingDraftId] = React.useState<string | null>(null);
+
+  const sessionKey = React.useMemo(() => {
+    return currentSessionId ?? (newSessionDraftOpen ? 'draft' : null);
+  }, [currentSessionId, newSessionDraftOpen]);
+
+  const drafts = React.useMemo(() => {
+    if (!sessionKey || !fileLabel) return [];
+    const sessionDrafts = allDrafts[sessionKey] ?? [];
+    return sessionDrafts.filter((draft) => draft.source === source && draft.fileLabel === fileLabel);
+  }, [allDrafts, fileLabel, sessionKey, source]);
+
+  const reset = React.useCallback(() => {
+    setSelection(null);
+    setCommentText('');
+    setEditingDraftId(null);
+  }, []);
+
+  const cancel = React.useCallback(() => {
+    reset();
+  }, [reset]);
+
+  const startEdit = React.useCallback((draft: InlineCommentDraft) => {
+    const draftRange = normalizeLineRange(fromDraftRange(draft));
+    setSelection(draftRange);
+    setCommentText(draft.text);
+    setEditingDraftId(draft.id);
+  }, [fromDraftRange]);
+
+  const deleteDraft = React.useCallback((draft: InlineCommentDraft) => {
+    removeDraft(draft.sessionKey, draft.id);
+    if (editingDraftId === draft.id) {
+      reset();
+    }
+  }, [editingDraftId, removeDraft, reset]);
+
+  const saveComment = React.useCallback((textToSave: string, rangeOverride?: TRange) => {
+    const targetRange = rangeOverride ?? selection;
+    const trimmedText = textToSave.trim();
+    if (!targetRange || !trimmedText || !fileLabel) return;
+
+    if (!sessionKey) {
+      toast.error('Select a session to save comment');
+      return;
+    }
+
+    const normalizedRange = normalizeLineRange(targetRange);
+    const normalizedStoreRange = normalizeStoreRange(toStoreRange(normalizedRange));
+    const code = getCodeForRange(normalizedRange);
+
+    if (editingDraftId) {
+      updateDraft(sessionKey, editingDraftId, {
+        fileLabel,
+        startLine: normalizedStoreRange.startLine,
+        endLine: normalizedStoreRange.endLine,
+        side: normalizedStoreRange.side,
+        code,
+        language,
+        text: trimmedText,
+      });
+    } else {
+      addDraft({
+        sessionKey,
+        source,
+        fileLabel,
+        startLine: normalizedStoreRange.startLine,
+        endLine: normalizedStoreRange.endLine,
+        side: normalizedStoreRange.side,
+        code,
+        language,
+        text: trimmedText,
+      });
+    }
+
+    reset();
+  }, [addDraft, editingDraftId, fileLabel, getCodeForRange, language, reset, selection, sessionKey, source, toStoreRange, updateDraft]);
+
+  return {
+    sessionKey,
+    drafts,
+    selection,
+    setSelection,
+    commentText,
+    setCommentText,
+    editingDraftId,
+    setEditingDraftId,
+    reset,
+    cancel,
+    startEdit,
+    deleteDraft,
+    saveComment,
+    fromDraftRange,
+  };
+}

packages/ui/src/components/views/FilesView.tsx

@@ -61,13 +61,10 @@ import { getLanguageFromExtension, getImageMimeType, isImageFile } from '@/lib/t
 import { useRuntimeAPIs } from '@/hooks/useRuntimeAPIs';
 import { EditorView } from '@codemirror/view';
 import { useThemeSystem } from '@/contexts/useThemeSystem';
-import { useSessionStore } from '@/stores/useSessionStore';
 import { useUIStore } from '@/stores/useUIStore';
 import { useFilesViewTabsStore } from '@/stores/useFilesViewTabsStore';
 import { useGitStatus } from '@/stores/useGitStore';
-import { useInlineCommentDraftStore } from '@/stores/useInlineCommentDraftStore';
-import { InlineCommentCard } from '@/components/comments/InlineCommentCard';
-import { InlineCommentInput } from '@/components/comments/InlineCommentInput';
+import { buildCodeMirrorCommentWidgets, normalizeLineRange, useInlineCommentController } from '@/components/comments';
 import { opencodeClient } from '@/lib/opencode/client';
 import { useDirectoryShowHidden } from '@/lib/directoryShowHidden';
 import { useFilesViewShowGitignored } from '@/lib/filesViewShowGitignored';
@@ -628,15 +625,8 @@ export const FilesView: React.FC<FilesViewProps> = ({ mode = 'full' }) => {
   const [isDragging, setIsDragging] = React.useState(false);
 
   // Session/config for sending comments
-  const currentSessionId = useSessionStore((state) => state.currentSessionId);
   const setMainTabGuard = useUIStore((state) => state.setMainTabGuard);
 
-  const addDraft = useInlineCommentDraftStore((s) => s.addDraft);
-  const updateDraft = useInlineCommentDraftStore((s) => s.updateDraft);
-  const removeDraft = useInlineCommentDraftStore((s) => s.removeDraft);
-  const allDrafts = useInlineCommentDraftStore((s) => s.drafts);
-  const [editingDraftId, setEditingDraftId] = React.useState<string | null>(null);
-
   // Global mouseup to end drag selection
   React.useEffect(() => {
     const handleGlobalMouseUp = () => {
@@ -648,14 +638,6 @@ export const FilesView: React.FC<FilesViewProps> = ({ mode = 'full' }) => {
     return () => document.removeEventListener('mouseup', handleGlobalMouseUp);
   }, []);
 
-  // Clear selection when file changes
-  React.useEffect(() => {
-    setLineSelection(null);
-    setMainTabGuard(null);
-    setDraftContent('');
-    setIsSaving(false);
-  }, [selectedFile?.path, setMainTabGuard]);
-
   React.useEffect(() => {
     return () => {
       if (copiedContentTimeoutRef.current !== null) {
@@ -667,25 +649,60 @@ export const FilesView: React.FC<FilesViewProps> = ({ mode = 'full' }) => {
     };
   }, []);
 
-  // Click outside to dismiss selection
+  // Extract selected code
+  const extractSelectedCode = React.useCallback((content: string, range: SelectedLineRange): string => {
+    const lines = content.split('\n');
+    const startLine = Math.max(1, range.start);
+    const endLine = Math.min(lines.length, range.end);
+    if (startLine > endLine) return '';
+    return lines.slice(startLine - 1, endLine).join('\n');
+  }, []);
+
+  const fileCommentController = useInlineCommentController<SelectedLineRange>({
+    source: 'file',
+    fileLabel: selectedFile?.path ?? null,
+    language: selectedFile?.path ? getLanguageFromExtension(selectedFile.path) || 'text' : 'text',
+    getCodeForRange: (range) => extractSelectedCode(fileContent, normalizeLineRange(range)),
+    toStoreRange: (range) => ({ startLine: range.start, endLine: range.end }),
+    fromDraftRange: (draft) => ({ start: draft.startLine, end: draft.endLine }),
+  });
+
+  const {
+    drafts: filesFileDrafts,
+    commentText,
+    editingDraftId,
+    setSelection: setCommentSelection,
+    saveComment,
+    cancel,
+    reset,
+    startEdit,
+    deleteDraft,
+  } = fileCommentController;
+
+  React.useEffect(() => {
+    setLineSelection(null);
+    reset();
+    setMainTabGuard(null);
+    setDraftContent('');
+    setIsSaving(false);
+  }, [selectedFile?.path, reset, setMainTabGuard]);
+
+  React.useEffect(() => {
+    setCommentSelection(lineSelection);
+  }, [lineSelection, setCommentSelection]);
+
   React.useEffect(() => {
     if (!lineSelection && !editingDraftId) return;
 
     const handleClickOutside = (e: MouseEvent) => {
       const target = e.target as HTMLElement;
 
-      // Check if click is inside comment UI
       if (target.closest('[data-comment-input="true"]') || target.closest('[data-comment-card="true"]')) return;
-
-      // Check if click is on CM gutter (only gutter should not dismiss)
       if (target.closest('.cm-gutterElement')) return;
-
-      // Check if click is inside toast (sonner)
       if (target.closest('[data-sonner-toast]') || target.closest('[data-sonner-toaster]')) return;
 
-      // Clicking anywhere else (including code content) dismisses selection
       setLineSelection(null);
-      setEditingDraftId(null);
+      cancel();
     };
 
     const timeoutId = setTimeout(() => {
@@ -696,49 +713,16 @@ export const FilesView: React.FC<FilesViewProps> = ({ mode = 'full' }) => {
       clearTimeout(timeoutId);
       document.removeEventListener('click', handleClickOutside);
     };
-  }, [lineSelection, editingDraftId]);
-
-  // Extract selected code
-  const extractSelectedCode = React.useCallback((content: string, range: SelectedLineRange): string => {
-    const lines = content.split('\n');
-    const startLine = Math.max(1, range.start);
-    const endLine = Math.min(lines.length, range.end);
-    if (startLine > endLine) return '';
-    return lines.slice(startLine - 1, endLine).join('\n');
-  }, []);
+  }, [cancel, editingDraftId, lineSelection]);
 
   const handleSaveComment = React.useCallback((text: string, range?: { start: number; end: number }) => {
-    if (!selectedFile) return;
-
-    const sessionKey = currentSessionId ?? 'draft';
-    const finalRange = range || lineSelection;
-    if (!finalRange) return;
-
-    const code = extractSelectedCode(fileContent, { start: finalRange.start, end: finalRange.end });
-
-    if (editingDraftId) {
-      updateDraft(sessionKey, editingDraftId, {
-        text: text.trim(),
-        code,
-        startLine: finalRange.start,
-        endLine: finalRange.end,
-      });
-    } else {
-      addDraft({
-        sessionKey,
-        source: 'file',
-        fileLabel: selectedFile.path,
-        startLine: finalRange.start,
-        endLine: finalRange.end,
-        code,
-        language: getLanguageFromExtension(selectedFile.path) || 'text',
-        text: text.trim(),
-      });
+    const finalRange = range ?? lineSelection ?? undefined;
+    if (range) {
+      setLineSelection(range);
     }
-
+    saveComment(text, finalRange);
     setLineSelection(null);
-    setEditingDraftId(null);
-  }, [selectedFile, currentSessionId, lineSelection, fileContent, extractSelectedCode, editingDraftId, updateDraft, addDraft]);
+  }, [lineSelection, saveComment]);
 
   const mapDirectoryEntries = React.useCallback((dirPath: string, entries: Array<{ name: string; path: string; isDirectory: boolean }>): FileNode[] => {
     const nodes = entries
@@ -1792,72 +1776,28 @@ export const FilesView: React.FC<FilesViewProps> = ({ mode = 'full' }) => {
       </Dialog>
     );
 
-  const filesFileDrafts = React.useMemo(() => {
-    if (!selectedFile) return [];
-    const sessionKey = currentSessionId ?? 'draft';
-    const sessionDrafts = allDrafts[sessionKey] ?? [];
-    return sessionDrafts.filter((d) => d.source === 'file' && d.fileLabel === selectedFile.path);
-  }, [selectedFile, currentSessionId, allDrafts]);
-
   const blockWidgets = React.useMemo(() => {
-    const widgets: import('@/components/ui/CodeMirrorEditor').BlockWidgetDef[] = [];
-
-    for (const draft of filesFileDrafts) {
-      if (draft.id === editingDraftId) {
-        widgets.push({
-          afterLine: draft.endLine,
-          id: `edit-${draft.id}`,
-          content: (
-            <InlineCommentInput
-              key={draft.id}
-              initialText={draft.text}
-              fileLabel={selectedFile?.path}
-              lineRange={{ start: draft.startLine, end: draft.endLine }}
-              isEditing={true}
-              onSave={handleSaveComment}
-              onCancel={() => setEditingDraftId(null)}
-            />
-          ),
-        });
-      } else {
-        widgets.push({
-          afterLine: draft.endLine,
-          id: `card-${draft.id}`,
-          content: (
-            <InlineCommentCard
-              key={draft.id}
-              draft={draft}
-              onEdit={() => {
-                setEditingDraftId(draft.id);
-                setLineSelection(null);
-              }}
-              onDelete={() => removeDraft(draft.sessionKey, draft.id)}
-            />
-          ),
-        });
-      }
-    }
-
-    if (lineSelection && !editingDraftId && !isDragging) {
-      widgets.push({
-        afterLine: lineSelection.end,
-        id: 'files-new-comment-input',
-        content: (
-          <InlineCommentInput
-            key="new-comment"
-            initialText=""
-            fileLabel={selectedFile?.path}
-            lineRange={lineSelection}
-            isEditing={false}
-            onSave={handleSaveComment}
-            onCancel={() => setLineSelection(null)}
-          />
-        ),
-      });
-    }
-
-    return widgets;
-  }, [filesFileDrafts, editingDraftId, lineSelection, isDragging, selectedFile?.path, handleSaveComment, removeDraft]);
+    return buildCodeMirrorCommentWidgets({
+      drafts: filesFileDrafts,
+      editingDraftId,
+      commentText,
+      selection: lineSelection,
+      isDragging,
+      fileLabel: selectedFile?.path ?? '',
+      newWidgetId: 'files-new-comment-input',
+      mapDraftToRange: (draft) => ({ start: draft.startLine, end: draft.endLine }),
+      onSave: handleSaveComment,
+      onCancel: () => {
+        setLineSelection(null);
+        cancel();
+      },
+      onEdit: (draft) => {
+        startEdit(draft);
+        setLineSelection({ start: draft.startLine, end: draft.endLine });
+      },
+      onDelete: deleteDraft,
+    });
+  }, [cancel, commentText, deleteDraft, editingDraftId, filesFileDrafts, handleSaveComment, isDragging, lineSelection, selectedFile?.path, startEdit]);
 
   const fileViewer = (
     <div

packages/ui/src/components/views/PierreDiffViewer.tsx

@@ -1,17 +1,22 @@
-import React, { useMemo, useRef, useState, useCallback, useEffect } from 'react';
-import { createPortal } from 'react-dom';
+import React, { useMemo, useRef, useCallback, useEffect } from 'react';
 import {
   FileDiff as PierreFileDiff,
   VirtualizedFileDiff,
   Virtualizer,
   type FileContents,
   type FileDiffOptions,
-  type SelectedLineRange,
   type DiffLineAnnotation,
+  type SelectedLineRange,
   type AnnotationSide,
   type VirtualFileMetrics,
 } from '@pierre/diffs';
-import { InlineCommentCard, InlineCommentInput } from '@/components/comments';
+import {
+  buildPierreLineAnnotations,
+  type PierreAnnotationData,
+  PierreDiffCommentOverlays,
+  toPierreAnnotationId,
+  useInlineCommentController,
+} from '@/components/comments';
 
 import { useOptionalThemeSystem } from '@/contexts/useThemeSystem';
 import { ScrollableOverlay } from '@/components/ui/ScrollableOverlay';
@@ -19,12 +24,9 @@ import { useWorkerPool } from '@/contexts/DiffWorkerProvider';
 import { ensurePierreThemeRegistered, getResolvedShikiTheme } from '@/lib/shiki/appThemeRegistry';
 import { getDefaultTheme } from '@/lib/theme/themes';
 
-import { toast } from '@/components/ui';
-import { useSessionStore } from '@/stores/useSessionStore';
 import { useUIStore } from '@/stores/useUIStore';
 import { useDeviceInfo } from '@/lib/device';
 import { cn } from '@/lib/utils';
-import { useInlineCommentDraftStore, type InlineCommentDraft } from '@/stores/useInlineCommentDraftStore';
 
 
 interface PierreDiffViewerProps {
@@ -113,10 +115,6 @@ const isSameSelection = (left: SelectedLineRange | null, right: SelectedLineRang
   return left.start === right.start && left.end === right.end && left.side === right.side;
 };
 
-type AnnotationData = 
-  | { type: 'saved' | 'edit'; draft: InlineCommentDraft }
-  | { type: 'new'; selection: SelectedLineRange };
-
 type SharedVirtualizer = {
   virtualizer: Virtualizer;
   release: () => void;
@@ -210,20 +208,37 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
 
   useUIStore();
   const { isMobile } = useDeviceInfo();
-  
-  const addDraft = useInlineCommentDraftStore((state) => state.addDraft);
-  const updateDraft = useInlineCommentDraftStore((state) => state.updateDraft);
-  const removeDraft = useInlineCommentDraftStore((state) => state.removeDraft);
-  const allDrafts = useInlineCommentDraftStore((state) => state.drafts);
-  const currentSessionId = useSessionStore((state) => state.currentSessionId);
 
-  const getSessionKey = useCallback(() => {
-    return currentSessionId ?? 'draft';
-  }, [currentSessionId]);
+  const diffCommentController = useInlineCommentController<SelectedLineRange>({
+    source: 'diff',
+    fileLabel: fileName || 'unknown',
+    language,
+    getCodeForRange: (range) => extractSelectedCode(original, modified, range),
+    toStoreRange: (range) => ({
+      startLine: range.start,
+      endLine: range.end,
+      side: range.side === 'deletions' ? 'original' : 'modified',
+    }),
+    fromDraftRange: (draft) => ({
+      start: draft.startLine,
+      end: draft.endLine,
+      side: draft.side === 'original' ? 'deletions' : 'additions',
+    }),
+  });
+
+  const {
+    drafts: fileDrafts,
+    selection,
+    setSelection,
+    commentText,
+    setCommentText,
+    editingDraftId,
+    saveComment,
+    cancel,
+    startEdit,
+    deleteDraft,
+  } = diffCommentController;
 
-  const [selection, setSelection] = useState<SelectedLineRange | null>(null);
-  const [commentText, setCommentText] = useState('');
-  const [editingDraftId, setEditingDraftId] = useState<string | null>(null);
   const selectionRef = useRef<SelectedLineRange | null>(null);
   const editingDraftIdRef = useRef<string | null>(null);
   // Use a ref to track if we're currently applying a selection programmatically
@@ -262,117 +277,27 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
         setCommentText('');
       }
     }
-  }, [isMobile]);
+  }, [isMobile, setCommentText, setSelection]);
 
   const handleCancelComment = useCallback(() => {
-    setCommentText('');
-    setSelection(null);
-    setEditingDraftId(null);
-  }, []);
+    cancel();
+  }, [cancel]);
 
-  // Helper to generate consistent annotation IDs
-  const getAnnotationId = useCallback((meta: AnnotationData): string => {
-    if (meta.type === 'saved' || meta.type === 'edit') {
-      return `draft-${meta.draft.id}`;
-    } else if (meta.type === 'new') {
-      const start = Math.min(meta.selection.start, meta.selection.end);
-      const end = Math.max(meta.selection.start, meta.selection.end);
-      const side = meta.selection.side ?? 'additions';
-      return `new-comment-${side}-${start}-${end}`;
-    }
-    return '';
-  }, []);
-
-  const renderAnnotation = useCallback((annotation: DiffLineAnnotation<AnnotationData>) => {
+  const renderAnnotation = useCallback((annotation: DiffLineAnnotation<PierreAnnotationData>) => {
     const div = document.createElement('div');
     div.style.position = 'relative';
 
-    const meta = (annotation as DiffLineAnnotation<AnnotationData>).metadata;
-    const id = getAnnotationId(meta);
+    const id = toPierreAnnotationId(annotation.metadata);
 
     div.dataset.annotationId = id;
     div.dataset.annotationSide = annotation.side;
     div.dataset.annotationLine = String(annotation.lineNumber);
     return div;
-  }, [getAnnotationId]);
-
-  const annotationTargetsRef = useRef<Record<string, HTMLElement | null>>({});
-
-  const resolveAnnotationTarget = useCallback((id: string): HTMLElement | null => {
-    const wrapper = diffRootRef.current;
-    if (!wrapper) return null;
-
-    const cached = annotationTargetsRef.current[id];
-    if (cached && wrapper.contains(cached)) {
-      return cached;
-    }
-
-    const host = wrapper.querySelector('diffs-container');
-    if (!host) return null;
-
-    const shadowRoot = host.shadowRoot;
-    if (!shadowRoot) return null;
-
-    const target = shadowRoot.querySelector(`[data-annotation-id="${id}"]`) as HTMLElement | null;
-    annotationTargetsRef.current[id] = target;
-    return target;
   }, []);
 
   const handleSaveComment = useCallback((textToSave: string, rangeOverride?: SelectedLineRange) => {
-    // Use provided range override or fall back to current selection
-    const targetRange = rangeOverride ?? selection;
-    if (!targetRange || !textToSave.trim()) return;
-
-    const normalizedStart = Math.min(targetRange.start, targetRange.end);
-    const normalizedEnd = Math.max(targetRange.start, targetRange.end);
-    const normalizedRange: SelectedLineRange = {
-      ...targetRange,
-      start: normalizedStart,
-      end: normalizedEnd,
-    };
-
-    const sessionKey = getSessionKey();
-    if (!sessionKey) {
-      toast.error('Select a session to save comment');
-      return;
-    }
-
-    // Pierre selection range: { start, end, side }
-    // Store needs { startLine, endLine, side: 'original'|'modified' }
-    // Pierre side: 'additions' (right) | 'deletions' (left)
-    const storeSide = normalizedRange.side === 'deletions' ? 'original' : 'modified';
-
-    // Use deterministic code extraction instead of instance.getSelectedText()
-    const selectedText = extractSelectedCode(original, modified, normalizedRange);
-
-    if (editingDraftId) {
-      updateDraft(sessionKey, editingDraftId, {
-        fileLabel: fileName || 'unknown',
-        startLine: normalizedRange.start,
-        endLine: normalizedRange.end,
-        side: storeSide,
-        code: selectedText,
-        language: language,
-        text: textToSave.trim(),
-      });
-    } else {
-      addDraft({
-        sessionKey,
-        source: 'diff',
-        fileLabel: fileName || 'unknown',
-        startLine: normalizedRange.start,
-        endLine: normalizedRange.end,
-        side: storeSide,
-        code: selectedText,
-        language: language,
-        text: textToSave.trim(),
-      });
-    }
-
-    setCommentText('');
-    setSelection(null);
-    setEditingDraftId(null);
-  }, [selection, fileName, language, original, modified, addDraft, updateDraft, getSessionKey, editingDraftId]);
+    saveComment(textToSave, rangeOverride ?? selection ?? undefined);
+  }, [saveComment, selection]);
 
 
   const applySelection = useCallback((range: SelectedLineRange) => {
@@ -388,6 +313,40 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
     } finally {
       isApplyingSelectionRef.current = false;
     }
+  }, [setSelection]);
+
+  const resolveClickedSide = useCallback((numberCell: HTMLElement): AnnotationSide => {
+    const lineType =
+      numberCell.closest('[data-line-type]')?.getAttribute('data-line-type')
+      ?? numberCell.getAttribute('data-line-type');
+    if (lineType === 'change-deletion') {
+      return 'deletions';
+    }
+    if (lineType === 'change-addition') {
+      return 'additions';
+    }
+
+    const explicitColumnSide =
+      numberCell.getAttribute('data-column-side')
+      ?? numberCell.getAttribute('data-side')
+      ?? numberCell.closest('[data-column-side]')?.getAttribute('data-column-side');
+    if (explicitColumnSide === 'deletions' || explicitColumnSide === 'left' || explicitColumnSide === 'original') {
+      return 'deletions';
+    }
+    if (explicitColumnSide === 'additions' || explicitColumnSide === 'right' || explicitColumnSide === 'modified') {
+      return 'additions';
+    }
+
+    const row = numberCell.closest('[data-line-type]');
+    if (row instanceof HTMLElement) {
+      const rowRect = row.getBoundingClientRect();
+      const cellRect = numberCell.getBoundingClientRect();
+      const rowCenter = rowRect.left + rowRect.width / 2;
+      const cellCenter = cellRect.left + cellRect.width / 2;
+      return cellCenter < rowCenter ? 'deletions' : 'additions';
+    }
+
+    return 'additions';
   }, []);
 
   ensurePierreThemeRegistered(lightTheme);
@@ -505,46 +464,12 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
 
 
   const lineAnnotations = useMemo(() => {
-    const sessionKey = getSessionKey();
-    if (!sessionKey) return [];
-
-    const sessionDrafts = allDrafts[sessionKey] ?? [];
-    // Match file label logic - use basename
-    const fileLabel = fileName || 'unknown';
-    const fileDrafts = sessionDrafts.filter((d) => d.source === 'diff' && d.fileLabel === fileLabel);
-
-    const anns: DiffLineAnnotation<AnnotationData>[] = [];
-
-    fileDrafts.forEach((d) => {
-      // Force cast to AnnotationSide to satisfy compiler
-      const side = (d.side === 'original' ? 'deletions' : 'additions') as AnnotationSide;
-      if (d.id === editingDraftId) {
-        // Always show edit input (even on mobile)
-        anns.push({
-          lineNumber: d.endLine,
-          side: side,
-          metadata: { type: 'edit', draft: d },
-        });
-      } else {
-        // Show saved cards on all devices
-        anns.push({
-          lineNumber: d.endLine,
-          side: side,
-          metadata: { type: 'saved', draft: d },
-        });
-      }
+    return buildPierreLineAnnotations({
+      drafts: fileDrafts,
+      editingDraftId,
+      selection,
     });
-
-    if (selection && !editingDraftId) {
-      anns.push({
-        lineNumber: selection.end,
-        side: (selection.side ?? 'additions') as AnnotationSide,
-        metadata: { type: 'new', selection },
-      });
-    }
-
-    return anns;
-  }, [allDrafts, getSessionKey, fileName, editingDraftId, selection]);
+  }, [editingDraftId, fileDrafts, selection]);
 
   const lineAnnotationsRef = useRef(lineAnnotations);
 
@@ -686,11 +611,7 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
         const lineNumber = lineRaw ? parseInt(lineRaw, 10) : NaN;
         if (Number.isNaN(lineNumber)) return;
 
-        const lineType =
-          numberCell.closest('[data-line-type]')?.getAttribute('data-line-type')
-          ?? numberCell.getAttribute('data-line-type');
-
-        const side: AnnotationSide = lineType === 'change-deletion' ? 'deletions' : 'additions';
+        const side = resolveClickedSide(numberCell);
 
         handleSelectionChange({
           start: lineNumber,
@@ -716,7 +637,7 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
       }
       cleanup();
     };
-  }, [diffThemeKey, fileName, handleSelectionChange]);
+  }, [diffThemeKey, fileName, handleSelectionChange, resolveClickedSide]);
 
   // MutationObserver to trigger re-renders when annotation DOM nodes are added/removed
   useEffect(() => {
@@ -728,10 +649,9 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
 
     const setupObserver = () => {
       const diffsContainer = container.querySelector('diffs-container');
-      if (!diffsContainer) return;
+      if (!(diffsContainer instanceof HTMLElement)) return;
 
       const shadowRoot = diffsContainer.shadowRoot;
-      if (!shadowRoot) return;
 
       observer = new MutationObserver(() => {
         if (rafId) cancelAnimationFrame(rafId);
@@ -741,7 +661,10 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
         });
       });
 
-      observer.observe(shadowRoot, { childList: true, subtree: true });
+      observer.observe(diffsContainer, { childList: true, subtree: true });
+      if (shadowRoot) {
+        observer.observe(shadowRoot, { childList: true, subtree: true });
+      }
     };
 
     const timeoutId = setTimeout(setupObserver, 100);
@@ -757,71 +680,26 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
     return null;
   }
 
-  const sessionKey = getSessionKey();
-  const sessionDrafts = sessionKey ? (allDrafts[sessionKey] ?? []) : [];
-  const fileLabel = fileName || 'unknown';
-  const fileDrafts = sessionDrafts.filter((d) => d.source === 'diff' && d.fileLabel === fileLabel);
-
-  const commentPortals = (
-    <>
-      {fileDrafts.map((d) => {
-        const target = resolveAnnotationTarget(`draft-${d.id}`);
-        if (!target) return null;
-
-        if (d.id === editingDraftId) {
-          return createPortal(
-            <InlineCommentInput
-              initialText={commentText}
-              fileLabel={(fileName?.split('/').pop()) ?? ''}
-              lineRange={{
-                start: d.startLine,
-                end: d.endLine,
-                side: d.side === 'original' ? 'deletions' : 'additions'
-              }}
-              isEditing={true}
-              onSave={handleSaveComment}
-              onCancel={handleCancelComment}
-            />,
-            target,
-            `draft-edit-${d.id}`
-          );
-        }
-
-        return createPortal(
-          <InlineCommentCard
-            draft={d}
-            onEdit={() => {
-              const side = d.side === 'original' ? 'deletions' : 'additions';
-              applySelection({ start: d.startLine, end: d.endLine, side });
-              setCommentText(d.text);
-              setEditingDraftId(d.id);
-            }}
-            onDelete={() => removeDraft(d.sessionKey, d.id)}
-          />,
-          target,
-          `draft-card-${d.id}`
-        );
-      })}
-
-      {selection && !editingDraftId && (() => {
-        const newCommentAnnotationId = getAnnotationId({ type: 'new', selection });
-        const target = resolveAnnotationTarget(newCommentAnnotationId);
-        if (!target) return null;
-
-        return createPortal(
-          <InlineCommentInput
-            initialText={commentText}
-            fileLabel={(fileName?.split('/').pop()) ?? ''}
-            lineRange={selection || undefined}
-            isEditing={false}
-            onSave={handleSaveComment}
-            onCancel={handleCancelComment}
-          />,
-          target,
-          newCommentAnnotationId
-        );
-      })()}
-    </>
+  const commentOverlays = (
+    <PierreDiffCommentOverlays
+      diffRootRef={diffRootRef}
+      drafts={fileDrafts}
+      selection={selection}
+      editingDraftId={editingDraftId}
+      commentText={commentText}
+      fileLabel={(fileName?.split('/').pop()) ?? ''}
+      onSave={handleSaveComment}
+      onCancel={handleCancelComment}
+      onEdit={(draft) => {
+        applySelection({
+          start: draft.startLine,
+          end: draft.endLine,
+          side: draft.side === 'original' ? 'deletions' : 'additions',
+        });
+        startEdit(draft);
+      }}
+      onDelete={deleteDraft}
+    />
   );
 
   if (layout === 'fill') {
@@ -838,7 +716,7 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
               <div ref={diffContainerRef} className="size-full" />
             </div>
           </ScrollableOverlay>
-          {commentPortals}
+          {commentOverlays}
         </div>
       </div>
     );
@@ -850,7 +728,7 @@ export const PierreDiffViewer: React.FC<PierreDiffViewerProps> = ({
       <div ref={diffRootRef} className="pierre-diff-wrapper w-full overflow-x-auto overflow-y-visible relative">
         <div ref={diffContainerRef} className="w-full" />
       </div>
-      {commentPortals}
+      {commentOverlays}
     </div>
   );
 };

packages/ui/src/components/views/PlanView.tsx

@@ -6,8 +6,7 @@ import { ErrorBoundary } from '@/components/ui/ErrorBoundary';
 import { ScrollableOverlay } from '@/components/ui/ScrollableOverlay';
 import { Button } from '@/components/ui/button';
 import { useUIStore } from '@/stores/useUIStore';
-import { InlineCommentCard } from '@/components/comments/InlineCommentCard';
-import { InlineCommentInput } from '@/components/comments/InlineCommentInput';
+import { buildCodeMirrorCommentWidgets, normalizeLineRange, useInlineCommentController } from '@/components/comments';
 
 import { getLanguageFromExtension } from '@/lib/toolHelpers';
 import { useDeviceInfo } from '@/lib/device';
@@ -20,8 +19,6 @@ import { useSessionStore } from '@/stores/useSessionStore';
 import { useDirectoryStore } from '@/stores/useDirectoryStore';
 import { useRuntimeAPIs } from '@/hooks/useRuntimeAPIs';
 import { EditorView } from '@codemirror/view';
-import { useInlineCommentDraftStore } from '@/stores/useInlineCommentDraftStore';
-import { toast } from '@/components/ui';
 import { copyTextToClipboard } from '@/lib/clipboard';
 
 const normalize = (value: string): string => {
@@ -87,23 +84,11 @@ export const PlanView: React.FC = () => {
   const sessions = useSessionStore((state) => state.sessions);
   const homeDirectory = useDirectoryStore((state) => state.homeDirectory);
   const runtimeApis = useRuntimeAPIs();
-  const newSessionDraftOpen = useSessionStore((state) => state.newSessionDraft?.open);
   useUIStore();
   const { isMobile } = useDeviceInfo();
   const { currentTheme } = useThemeSystem();
   React.useMemo(() => generateSyntaxTheme(currentTheme), [currentTheme]);
 
-  // Inline comment drafts
-  const addDraft = useInlineCommentDraftStore((state) => state.addDraft);
-  const updateDraft = useInlineCommentDraftStore((state) => state.updateDraft);
-  const removeDraft = useInlineCommentDraftStore((state) => state.removeDraft);
-  const allDrafts = useInlineCommentDraftStore((state) => state.drafts);
-
-  // Get session key for drafts
-  const getSessionKey = React.useCallback(() => {
-    return currentSessionId ?? (newSessionDraftOpen ? 'draft' : null);
-  }, [currentSessionId, newSessionDraftOpen]);
-
   const session = React.useMemo(() => {
     if (!currentSessionId) return null;
     return sessions.find((s) => s.id === currentSessionId) ?? null;
@@ -122,6 +107,9 @@ export const PlanView: React.FC = () => {
     return toDisplayPath(resolvedPath, { currentDirectory: sessionDirectory, homeDirectory });
   }, [resolvedPath, sessionDirectory, homeDirectory]);
   const [content, setContent] = React.useState<string>('');
+  const planFileLabel = React.useMemo(() => {
+    return displayPath ? displayPath.split('/').pop() || 'plan' : 'plan';
+  }, [displayPath]);
   const [loading, setLoading] = React.useState(false);
   const [copiedPath, setCopiedPath] = React.useState(false);
   const [copiedContent, setCopiedContent] = React.useState(false);
@@ -130,8 +118,6 @@ export const PlanView: React.FC = () => {
   const copiedContentTimeoutRef = React.useRef<number | null>(null);
 
   const [lineSelection, setLineSelection] = React.useState<SelectedLineRange | null>(null);
-  const [commentText, setCommentText] = React.useState('');
-  const [editingDraftId, setEditingDraftId] = React.useState<string | null>(null);
   const editorViewRef = React.useRef<EditorView | null>(null);
   const editorWrapperRef = React.useRef<HTMLDivElement | null>(null);
 
@@ -172,23 +158,66 @@ export const PlanView: React.FC = () => {
     return () => document.removeEventListener('mouseup', handleGlobalMouseUp);
   }, []);
 
+  const extractSelectedCode = React.useCallback((text: string, range: SelectedLineRange): string => {
+    const lines = text.split('\n');
+    const startLine = Math.max(1, range.start);
+    const endLine = Math.min(lines.length, range.end);
+    if (startLine > endLine) return '';
+    return lines.slice(startLine - 1, endLine).join('\n');
+  }, []);
+
+  const commentController = useInlineCommentController<SelectedLineRange>({
+    source: 'plan',
+    fileLabel: planFileLabel,
+    language: resolvedPath ? getLanguageFromExtension(resolvedPath) || 'markdown' : 'markdown',
+    getCodeForRange: (range) => extractSelectedCode(content, normalizeLineRange(range)),
+    toStoreRange: (range) => ({ startLine: range.start, endLine: range.end }),
+    fromDraftRange: (draft) => ({ start: draft.startLine, end: draft.endLine }),
+  });
+
+  const {
+    drafts: planFileDrafts,
+    commentText,
+    editingDraftId,
+    setSelection: setCommentSelection,
+    saveComment,
+    cancel,
+    reset,
+    startEdit,
+    deleteDraft,
+  } = commentController;
+
   React.useEffect(() => {
     setLineSelection(null);
-    setCommentText('');
-    setEditingDraftId(null);
-  }, [content]);
+    reset();
+  }, [content, reset]);
+
+  React.useEffect(() => {
+    setCommentSelection(lineSelection);
+  }, [lineSelection, setCommentSelection]);
+
+  const handleCancelComment = React.useCallback(() => {
+    setLineSelection(null);
+    cancel();
+  }, [cancel]);
+
+  const handleSaveComment = React.useCallback((textToSave: string, rangeOverride?: { start: number; end: number }) => {
+    if (rangeOverride) {
+      setLineSelection(rangeOverride);
+    }
+    saveComment(textToSave, rangeOverride ?? lineSelection ?? undefined);
+    setLineSelection(null);
+  }, [lineSelection, saveComment]);
 
   React.useEffect(() => {
     if (!lineSelection) return;
-    
-    // Auto-scroll input into view on mobile
+
     if (isMobile && !editingDraftId) {
-       // We rely on InlineCommentInput doing this now via useEffect
+      // Input handles mobile scroll/focus behavior.
     }
 
     const handleClickOutside = (e: MouseEvent) => {
       const target = e.target as HTMLElement;
-      // Check if click is inside any comment component
       if (
         target.closest('[data-comment-card="true"]') ||
         target.closest('[data-comment-input="true"]') ||
@@ -196,15 +225,12 @@ export const PlanView: React.FC = () => {
       ) {
         return;
       }
-      
+
       if (target.closest('.cm-gutterElement')) return;
       if (target.closest('[data-sonner-toast]') || target.closest('[data-sonner-toaster]')) return;
-      
-      // If clicking outside while editing, maybe we should save or ask confirmation?
-      // For now, cancel edit
+
       setLineSelection(null);
-      setCommentText('');
-      setEditingDraftId(null);
+      cancel();
     };
 
     const timeoutId = window.setTimeout(() => {
@@ -215,63 +241,7 @@ export const PlanView: React.FC = () => {
       window.clearTimeout(timeoutId);
       document.removeEventListener('click', handleClickOutside);
     };
-  }, [lineSelection, editingDraftId, isMobile]);
-
-  const extractSelectedCode = React.useCallback((text: string, range: SelectedLineRange): string => {
-    const lines = text.split('\n');
-    const startLine = Math.max(1, range.start);
-    const endLine = Math.min(lines.length, range.end);
-    if (startLine > endLine) return '';
-    return lines.slice(startLine - 1, endLine).join('\n');
-  }, []);
-
-  const handleCancelComment = React.useCallback(() => {
-    setCommentText('');
-    setLineSelection(null);
-    setEditingDraftId(null);
-  }, []);
-
-  const handleSaveComment = React.useCallback((textToSave: string, rangeOverride?: { start: number; end: number }) => {
-    // Use provided range override or fall back to current selection
-    const targetRange = rangeOverride ?? lineSelection;
-    if (!targetRange || !textToSave.trim()) return;
-
-    const sessionKey = getSessionKey();
-    if (!sessionKey) {
-      toast.error('Select a session to save comment');
-      return;
-    }
-
-    const code = extractSelectedCode(content, targetRange);
-    const fileLabel = displayPath ? displayPath.split('/').pop() || 'plan' : 'plan';
-    const language = resolvedPath ? getLanguageFromExtension(resolvedPath) || 'markdown' : 'markdown';
-
-    if (editingDraftId) {
-      updateDraft(sessionKey, editingDraftId, {
-        fileLabel,
-        startLine: targetRange.start,
-        endLine: targetRange.end,
-        code,
-        language,
-        text: textToSave.trim(),
-      });
-    } else {
-      addDraft({
-        sessionKey,
-        source: 'plan',
-        fileLabel,
-        startLine: targetRange.start,
-        endLine: targetRange.end,
-        code,
-        language,
-        text: textToSave.trim(),
-      });
-    }
-
-    setCommentText('');
-    setLineSelection(null);
-    setEditingDraftId(null);
-  }, [lineSelection, content, displayPath, resolvedPath, addDraft, updateDraft, getSessionKey, extractSelectedCode, editingDraftId]);
+  }, [cancel, editingDraftId, isMobile, lineSelection]);
 
 
   const editorExtensions = React.useMemo(() => {
@@ -381,64 +351,25 @@ export const PlanView: React.FC = () => {
     };
   }, []);
 
-  const planFileLabel = displayPath ? displayPath.split('/').pop() || 'plan' : 'plan';
-
-  const planFileDrafts = React.useMemo(() => {
-    const sessionKey = getSessionKey();
-    if (!sessionKey) return [];
-    const sessionDrafts = allDrafts[sessionKey] ?? [];
-    return sessionDrafts.filter((d) => d.source === 'plan' && d.fileLabel === planFileLabel);
-  }, [getSessionKey, allDrafts, planFileLabel]);
-
   const blockWidgets = React.useMemo(() => {
-    const widgets: Array<{ afterLine: number; id: string; content: React.ReactNode }> = [];
-
-    for (const draft of planFileDrafts) {
-      const content = draft.id === editingDraftId ? (
-        <InlineCommentInput
-          key={`edit-${draft.id}`}
-          initialText={commentText}
-          fileLabel={draft.fileLabel}
-          lineRange={{ start: draft.startLine, end: draft.endLine }}
-          isEditing={true}
-          onSave={handleSaveComment}
-          onCancel={handleCancelComment}
-        />
-      ) : (
-        <InlineCommentCard
-          key={`saved-${draft.id}`}
-          draft={draft}
-          onEdit={() => {
-            setLineSelection({ start: draft.startLine, end: draft.endLine });
-            setCommentText(draft.text);
-            setEditingDraftId(draft.id);
-          }}
-          onDelete={() => removeDraft(draft.sessionKey, draft.id)}
-        />
-      );
-      widgets.push({ afterLine: draft.endLine, id: draft.id, content });
-    }
-
-    if (lineSelection && !editingDraftId && !isDragging) {
-      widgets.push({
-        afterLine: Math.max(lineSelection.start, lineSelection.end),
-        id: 'plan-new-comment-input',
-        content: (
-          <InlineCommentInput
-            key="new-comment"
-            initialText={commentText}
-            fileLabel={planFileLabel}
-            lineRange={lineSelection}
-            isEditing={false}
-            onSave={handleSaveComment}
-            onCancel={handleCancelComment}
-          />
-        ),
-      });
-    }
-
-    return widgets;
-  }, [planFileDrafts, editingDraftId, commentText, lineSelection, isDragging, planFileLabel, handleSaveComment, handleCancelComment, removeDraft]);
+    return buildCodeMirrorCommentWidgets({
+      drafts: planFileDrafts,
+      editingDraftId,
+      commentText,
+      selection: lineSelection,
+      isDragging,
+      fileLabel: planFileLabel,
+      newWidgetId: 'plan-new-comment-input',
+      mapDraftToRange: (draft) => ({ start: draft.startLine, end: draft.endLine }),
+      onSave: handleSaveComment,
+      onCancel: handleCancelComment,
+      onEdit: (draft) => {
+        startEdit(draft);
+        setLineSelection({ start: draft.startLine, end: draft.endLine });
+      },
+      onDelete: deleteDraft,
+    });
+  }, [commentText, deleteDraft, editingDraftId, handleCancelComment, handleSaveComment, isDragging, lineSelection, planFileDrafts, planFileLabel, startEdit]);
 
   return (
     <div className="relative flex h-full min-h-0 min-w-0 w-full flex-col overflow-hidden bg-background">