fix: maintain formatting and add more context like selection

examples/superdoc-ai-quickstart/package.json

@@ -9,7 +9,7 @@
     "preview": "vite preview"
   },
   "dependencies": {
-    "@superdoc-dev/ai": "^0.1.3",
+    "@superdoc-dev/ai": "latest",
     "superdoc": "0.28.0"
   },
   "devDependencies": {

packages/ai/package.json

@@ -25,7 +25,7 @@
     "superdoc": "*"
   },
   "devDependencies": {
-    "superdoc": "^0.28.0",
+    "superdoc": "^0.29.0",
     "@types/node": "^20.0.0",
     "typescript": "^5.0.0",
     "eslint": "^8.0.0",

packages/ai/src/ai-actions-service.ts

@@ -15,6 +15,8 @@ import {
  */
 export class AIActionsService {
     private adapter: EditorAdapter;
+    private capturedContext: string | null = null;
+    private capturedSelectionBounds: { from: number; to: number } | null = null;
 
     constructor(
         private provider: AIProvider,
@@ -34,7 +36,30 @@ export class AIActionsService {
         }
     }
 
+    /**
+     * Sets a captured context that will be used instead of calling the provider.
+     * This ensures the context (including selection) is captured before async operations.
+     */
+    public setCapturedContext(context: string | null, selectionBounds?: { from: number; to: number } | null): void {
+        this.capturedContext = context;
+        this.capturedSelectionBounds = selectionBounds || null;
+    }
+
+    /**
+     * Clears the captured context, reverting to using the provider function.
+     */
+    public clearCapturedContext(): void {
+        this.capturedContext = null;
+        this.capturedSelectionBounds = null;
+    }
+
     private getDocumentContext(): string {
+        // If a context was captured synchronously, use it
+        if (this.capturedContext !== null) {
+            return this.capturedContext;
+        }
+
+        // Otherwise, call the provider function
         if (!this.documentContextProvider) {
             return '';
         }
@@ -80,7 +105,7 @@ export class AIActionsService {
         if (!result.success || !result.results) {
             return result;
         }
-        result.results = this.adapter.findResults(result.results);
+        result.results = this.adapter.findResults(result.results, this.capturedSelectionBounds);
 
         return result;
     }
@@ -184,7 +209,7 @@ export class AIActionsService {
             return [];
         }
 
-        const searchResults = this.adapter.findResults(replacements);
+        const searchResults = this.adapter.findResults(replacements, this.capturedSelectionBounds);
         const match = searchResults?.[0];
         for (const result of searchResults) {
             try {

packages/ai/src/ai-actions.ts

@@ -168,7 +168,35 @@ export class AIActions {
     }
 
     /**
-     * Executes an action with full callback lifecycle support
+     * Gets the current selection bounds if a selection exists.
+     * @private
+     * @returns Selection bounds {from, to} or null if no selection
+     */
+    private getSelectionBounds(): { from: number; to: number } | null {
+        const editor = this.getEditor();
+        if (!editor) {
+            return null;
+        }
+
+        const state = editor.view?.state || editor.state;
+        if (!state || !state.selection) {
+            return null;
+        }
+
+        const { selection } = state;
+        if (selection.empty) {
+            return null;
+        }
+
+        return {
+            from: selection.from,
+            to: selection.to,
+        };
+    }
+
+    /**
+     * Executes an action with full callback lifecycle support.
+     * Captures the document context (including selection) synchronously before any async operations.
      * @private
      */
     private async executeActionWithCallbacks<T extends Result>(
@@ -178,6 +206,13 @@ export class AIActions {
         if (!editor) {
             throw new Error('No active SuperDoc editor available for AI actions');
         }
+
+        // Capture context synchronously before any async operations
+        // This ensures the selection is locked in at the moment the action is called
+        const capturedContext = this.getDocumentContext();
+        const selectionBounds = this.getSelectionBounds();
+        this.commands.setCapturedContext(capturedContext, selectionBounds);
+
         try {
             this.callbacks.onStreamingStart?.();
             const result: T = await fn();
@@ -187,6 +222,9 @@ export class AIActions {
         } catch (error: Error | any) {
             this.handleError(error as Error);
             throw error;
+        } finally {
+            // Clear the captured context after the action completes
+            this.commands.clearCapturedContext();
         }
     }
 
@@ -302,17 +340,35 @@ export class AIActions {
 
     /**
      * Retrieves the current document context for AI processing.
-     * Combines XML and plain text representations when available.
+     * Returns selected text if available, otherwise returns the full document.
      *
-     * @returns Document context string
+     * @returns Document context string (selected text if available, otherwise full document)
      */
     public getDocumentContext(): string {
         const editor = this.getEditor();
         if (!editor) {
             return '';
         }
 
-        return editor.state?.doc?.textContent?.trim() || '';
+        // Try to get state from view first (most up-to-date), then fall back to editor.state
+        const state = editor.view?.state || editor.state;
+        if (!state || !state.doc) {
+            return '';
+        }
+
+        const { selection, doc } = state;
+
+        // If there's a non-empty selection, return the selected text
+        if (selection && !selection.empty) {
+            const selectedText = doc.textBetween(selection.from, selection.to, ' ').trim();
+            // Only return selected text if it's not empty (handles edge cases)
+            if (selectedText) {
+                return selectedText;
+            }
+        }
+
+        // Otherwise, return the full document content
+        return doc.textContent?.trim() || '';
     }
 
     /**

packages/ai/src/editor-adapter.ts

@@ -10,7 +10,8 @@ export class EditorAdapter {
     constructor(private editor: Editor) {}
 
     // Search for string occurrences and resolve document positions
-    findResults(results: FoundMatch[]): FoundMatch[] {
+    // If selectionBounds is provided, only returns matches within the selected area
+    findResults(results: FoundMatch[], selectionBounds?: { from: number; to: number } | null): FoundMatch[] {
         if (!results?.length) {
             return [];
         }
@@ -20,7 +21,7 @@ export class EditorAdapter {
                 const text = match.originalText;
                 const rawMatches = this.editor.commands?.search?.(text) ?? [];
 
-                const positions = rawMatches
+                let positions = rawMatches
                     .map((match: { from?: number; to?: number}) => {
                         const from = match.from;
                         const to = match.to;
@@ -29,7 +30,16 @@ export class EditorAdapter {
                         }
                         return { from, to };
                     })
-                .filter((value: { from: number; to: number } | null) => value !== null);
+                    .filter((value: { from: number; to: number } | null) => value !== null) as { from: number; to: number }[];
+
+                // Filter positions to only include those within the selection bounds if provided
+                if (selectionBounds) {
+                    positions = positions.filter((pos) => {
+                        // Check if the match overlaps with or is within the selection bounds
+                        // A match is within bounds if it starts at or after selection.from and ends at or before selection.to
+                        return pos.from >= selectionBounds.from && pos.to <= selectionBounds.to;
+                    });
+                }
 
                 return {
                     ...match,