diff --git a/frontend/appflowy_flutter/assets/test/workspaces/ai_workspace.zip b/frontend/appflowy_flutter/assets/test/workspaces/ai_workspace.zip
new file mode 100644
index 0000000000..b523b64ed0
Binary files /dev/null and b/frontend/appflowy_flutter/assets/test/workspaces/ai_workspace.zip differ
diff --git a/frontend/appflowy_flutter/integration_test/open_ai_smart_menu_test.dart b/frontend/appflowy_flutter/integration_test/open_ai_smart_menu_test.dart
new file mode 100644
index 0000000000..f63197a039
--- /dev/null
+++ b/frontend/appflowy_flutter/integration_test/open_ai_smart_menu_test.dart
@@ -0,0 +1,108 @@
+import 'package:appflowy/plugins/document/presentation/plugins/openai/service/openai_client.dart';
+import 'package:appflowy_editor/appflowy_editor.dart';
+import 'package:flutter/services.dart';
+import 'package:flutter_test/flutter_test.dart';
+import 'package:integration_test/integration_test.dart';
+import 'util/mock/mock_openai_repository.dart';
+import 'util/util.dart';
+import 'package:flowy_infra_ui/flowy_infra_ui.dart';
+import 'package:appflowy_editor/src/render/toolbar/toolbar_widget.dart';
+import 'package:appflowy/startup/startup.dart';
+
+void main() {
+  IntegrationTestWidgetsFlutterBinding.ensureInitialized();
+  const service = TestWorkspaceService(TestWorkspace.aiWorkSpace);
+
+  group('integration tests for open-ai smart menu', () {
+    setUpAll(() async => await service.setUpAll());
+    setUp(() async => await service.setUp());
+
+    testWidgets('testing selection on open-ai smart menu replace', (tester) async {
+      final appFlowyEditor = await setUpOpenAITesting(tester);
+      final editorState = appFlowyEditor.editorState;
+
+      editorState.service.selectionService.updateSelection(
+        Selection(
+          start: Position(path: [1], offset: 4),
+          end: Position(path: [1], offset: 10),
+        ),
+      );
+      await tester.pumpAndSettle(const Duration(milliseconds: 500));
+      await tester.pumpAndSettle();
+
+      expect(find.byType(ToolbarWidget), findsAtLeastNWidgets(1));
+
+      await tester.tap(find.byTooltip('AI Assistants'));
+      await tester.pumpAndSettle(const Duration(milliseconds: 500));
+
+      await tester.tap(find.text('Summarize'));
+      await tester.pumpAndSettle();
+
+      await tester.tap(find.byType(FlowyRichTextButton, skipOffstage: false).first);
+      await tester.pumpAndSettle();
+
+      expect(
+        editorState.service.selectionService.currentSelection.value,
+        Selection(
+          start: Position(path: [1], offset: 4),
+          end: Position(path: [1], offset: 84),
+        ),
+      );
+    });
+    testWidgets('testing selection on open-ai smart menu insert', (tester) async {
+      final appFlowyEditor = await setUpOpenAITesting(tester);
+      final editorState = appFlowyEditor.editorState;
+
+      editorState.service.selectionService.updateSelection(
+        Selection(
+          start: Position(path: [1], offset: 0),
+          end: Position(path: [1], offset: 5),
+        ),
+      );
+      await tester.pumpAndSettle(const Duration(milliseconds: 500));
+      await tester.pumpAndSettle();
+      expect(find.byType(ToolbarWidget), findsAtLeastNWidgets(1));
+
+      await tester.tap(find.byTooltip('AI Assistants'));
+      await tester.pumpAndSettle(const Duration(milliseconds: 500));
+
+      await tester.tap(find.text('Summarize'));
+      await tester.pumpAndSettle();
+
+      await tester.tap(find.byType(FlowyRichTextButton, skipOffstage: false).at(1));
+      await tester.pumpAndSettle();
+
+      expect(
+        editorState.service.selectionService.currentSelection.value,
+        Selection(
+          start: Position(path: [2], offset: 0),
+          end: Position(path: [3], offset: 0),
+        ),
+      );
+    });
+  });
+}
+
+Future<AppFlowyEditor> setUpOpenAITesting(WidgetTester tester) async {
+  await tester.initializeAppFlowy();
+  await mockOpenAIRepository();
+
+  await simulateKeyDownEvent(LogicalKeyboardKey.controlLeft);
+  await simulateKeyDownEvent(LogicalKeyboardKey.backslash);
+  await tester.pumpAndSettle();
+
+  final Finder editor = find.byType(AppFlowyEditor);
+  await tester.tap(editor);
+  await tester.pumpAndSettle();
+  return (tester.state(editor).widget as AppFlowyEditor);
+}
+
+Future<void> mockOpenAIRepository() async {
+  await getIt.unregister<OpenAIRepository>();
+  getIt.registerFactoryAsync<OpenAIRepository>(
+    () => Future.value(
+      MockOpenAIRepository(),
+    ),
+  );
+  return;
+}
diff --git a/frontend/appflowy_flutter/integration_test/runner.dart b/frontend/appflowy_flutter/integration_test/runner.dart
index 1592cccb93..b848ce1fc3 100644
--- a/frontend/appflowy_flutter/integration_test/runner.dart
+++ b/frontend/appflowy_flutter/integration_test/runner.dart
@@ -3,6 +3,7 @@ import 'package:integration_test/integration_test.dart';
 import 'board_test.dart' as board_test;
 import 'switch_folder_test.dart' as switch_folder_test;
 import 'empty_document_test.dart' as empty_document_test;
+import 'open_ai_smart_menu_test.dart' as smart_menu_test;
 
 /// The main task runner for all integration tests in AppFlowy.
 ///
@@ -16,4 +17,5 @@ void main() {
   switch_folder_test.main();
   board_test.main();
   empty_document_test.main();
+  smart_menu_test.main();
 }
diff --git a/frontend/appflowy_flutter/integration_test/util/data.dart b/frontend/appflowy_flutter/integration_test/util/data.dart
index c6912ddeca..0dd0961ee2 100644
--- a/frontend/appflowy_flutter/integration_test/util/data.dart
+++ b/frontend/appflowy_flutter/integration_test/util/data.dart
@@ -10,6 +10,7 @@ import 'package:shared_preferences/shared_preferences.dart';
 enum TestWorkspace {
   board("board"),
   emptyDocument("empty_document"),
+  aiWorkSpace("ai_workspace"),
   coverImage("cover_image");
 
   const TestWorkspace(this._name);
diff --git a/frontend/appflowy_flutter/integration_test/util/mock/mock_openai_repository.dart b/frontend/appflowy_flutter/integration_test/util/mock/mock_openai_repository.dart
new file mode 100644
index 0000000000..a9486a8536
--- /dev/null
+++ b/frontend/appflowy_flutter/integration_test/util/mock/mock_openai_repository.dart
@@ -0,0 +1,76 @@
+import 'package:appflowy/plugins/document/presentation/plugins/openai/service/openai_client.dart';
+import 'package:mocktail/mocktail.dart';
+import 'dart:convert';
+import 'package:appflowy/plugins/document/presentation/plugins/openai/service/text_completion.dart';
+import 'package:appflowy/plugins/document/presentation/plugins/openai/service/error.dart';
+import 'package:http/http.dart' as http;
+import 'dart:async';
+
+class MyMockClient extends Mock implements http.Client {
+  @override
+  Future<http.StreamedResponse> send(http.BaseRequest request) async {
+    final requestType = request.method;
+    final requestUri = request.url;
+
+    if (requestType == 'POST' && requestUri == OpenAIRequestType.textCompletion.uri) {
+      final responseHeaders = <String, String>{'content-type': 'text/event-stream'};
+      final responseBody = Stream.fromIterable([
+        utf8.encode(
+          '{ "choices": [{"text": "Lorem ipsum dolor sit amet, consectetuer adipiscing elit. Aenean commodo ligula ", "index": 0, "logprobs": null, "finish_reason": null}]}',
+        ),
+        utf8.encode('\n'),
+        utf8.encode('[DONE]'),
+      ]);
+
+      // Return a mocked response with the expected data
+      return http.StreamedResponse(responseBody, 200, headers: responseHeaders);
+    }
+
+    // Return an error response for any other request
+    return http.StreamedResponse(const Stream.empty(), 404);
+  }
+}
+
+class MockOpenAIRepository extends HttpOpenAIRepository {
+  MockOpenAIRepository() : super(apiKey: 'dummyKey', client: MyMockClient());
+
+  @override
+  Future<void> getStreamedCompletions({
+    required String prompt,
+    required Future<void> Function() onStart,
+    required Future<void> Function(TextCompletionResponse response) onProcess,
+    required Future<void> Function() onEnd,
+    required void Function(OpenAIError error) onError,
+    String? suffix,
+    int maxTokens = 2048,
+    double temperature = 0.3,
+    bool useAction = false,
+  }) async {
+    final request = http.Request('POST', OpenAIRequestType.textCompletion.uri);
+    final response = await client.send(request);
+
+    var previousSyntax = '';
+    if (response.statusCode == 200) {
+      await for (final chunk in response.stream.transform(const Utf8Decoder()).transform(const LineSplitter())) {
+        await onStart();
+        final data = chunk.trim().split('data: ');
+        if (data[0] != '[DONE]') {
+          final response = TextCompletionResponse.fromJson(
+            json.decode(data[0]),
+          );
+          if (response.choices.isNotEmpty) {
+            final text = response.choices.first.text;
+            if (text == previousSyntax && text == '\n') {
+              continue;
+            }
+            await onProcess(response);
+            previousSyntax = response.choices.first.text;
+          }
+        } else {
+          await onEnd();
+        }
+      }
+    }
+    return;
+  }
+}
diff --git a/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/service/openai_client.dart b/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/service/openai_client.dart
index ccb4b08866..5daf577564 100644
--- a/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/service/openai_client.dart
+++ b/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/service/openai_client.dart
@@ -50,6 +50,7 @@ abstract class OpenAIRepository {
     String? suffix,
     int maxTokens = 2048,
     double temperature = 0.3,
+    bool useAction = false,
   });
 
   ///  Get edits from GPT-3
diff --git a/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/widgets/smart_edit_node_widget.dart b/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/widgets/smart_edit_node_widget.dart
index 2c157f6ddc..6f580e7e7a 100644
--- a/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/widgets/smart_edit_node_widget.dart
+++ b/frontend/appflowy_flutter/lib/plugins/document/presentation/plugins/openai/widgets/smart_edit_node_widget.dart
@@ -4,7 +4,7 @@ import 'package:appflowy/plugins/document/presentation/plugins/openai/service/op
 import 'package:appflowy/plugins/document/presentation/plugins/openai/util/learn_more_action.dart';
 import 'package:appflowy/plugins/document/presentation/plugins/openai/widgets/discard_dialog.dart';
 import 'package:appflowy/plugins/document/presentation/plugins/openai/widgets/smart_edit_action.dart';
-import 'package:appflowy/user/application/user_service.dart';
+import 'package:appflowy/startup/startup.dart';
 import 'package:appflowy_editor/appflowy_editor.dart';
 import 'package:appflowy_popover/appflowy_popover.dart';
 import 'package:flowy_infra_ui/flowy_infra_ui.dart';
@@ -242,7 +242,7 @@ class _SmartEditInputState extends State<_SmartEditInput> {
           ),
           onPressed: () async {
             await _onReplace();
-            _onExit();
+            await _onExit();
           },
         ),
         const Space(10, 0),
@@ -257,7 +257,7 @@ class _SmartEditInputState extends State<_SmartEditInput> {
           ),
           onPressed: () async {
             await _onInsertBelow();
-            _onExit();
+            await _onExit();
           },
         ),
         const Space(10, 0),
@@ -272,10 +272,13 @@ class _SmartEditInputState extends State<_SmartEditInput> {
           ),
           onPressed: () async => await _onExit(),
         ),
-        const Spacer(),
-        FlowyText.regular(
-          LocaleKeys.document_plugins_warning.tr(),
-          color: Theme.of(context).hintColor,
+        const Spacer(flex: 2),
+        Expanded(
+          child: FlowyText.regular(
+            overflow: TextOverflow.ellipsis,
+            LocaleKeys.document_plugins_warning.tr(),
+            color: Theme.of(context).hintColor,
+          ),
         ),
       ],
     );
@@ -298,7 +301,22 @@ class _SmartEditInputState extends State<_SmartEditInput> {
       selection,
       texts,
     );
-    return widget.editorState.apply(transaction);
+    await widget.editorState.apply(transaction);
+
+    int endOffset = texts.last.length;
+    if (texts.length == 1) {
+      endOffset += selection.start.offset;
+    }
+
+    await widget.editorState.updateCursorSelection(
+      Selection(
+        start: selection.start,
+        end: Position(
+          path: [selection.start.path.first + texts.length - 1],
+          offset: endOffset,
+        ),
+      ),
+    );
   }
 
   Future<void> _onInsertBelow() async {
@@ -317,7 +335,16 @@ class _SmartEditInputState extends State<_SmartEditInput> {
         ),
       ),
     );
-    return widget.editorState.apply(transaction);
+    await widget.editorState.apply(transaction);
+
+    await widget.editorState.updateCursorSelection(
+      Selection(
+        start: Position(path: selection.end.path.next, offset: 0),
+        end: Position(
+          path: [selection.end.path.next.first + texts.length],
+        ),
+      ),
+    );
   }
 
   Future<void> _onExit() async {
@@ -333,51 +360,42 @@ class _SmartEditInputState extends State<_SmartEditInput> {
   }
 
   Future<void> _requestCompletions() async {
-    final result = await UserBackendService.getCurrentUserProfile();
-    return result.fold((l) async {
-      final openAIRepository = HttpOpenAIRepository(
-        client: client,
-        apiKey: l.openaiKey,
-      );
+    final openAIRepository = await getIt.getAsync<OpenAIRepository>();
 
-      var lines = input.split('\n\n');
-      if (action == SmartEditAction.summarize) {
-        lines = [lines.join('\n')];
-      }
-      for (var i = 0; i < lines.length; i++) {
-        final element = lines[i];
-        await openAIRepository.getStreamedCompletions(
-          useAction: true,
-          prompt: action.prompt(element),
-          onStart: () async {
-            setState(() {
-              loading = false;
-            });
-          },
-          onProcess: (response) async {
-            setState(() {
-              if (response.choices.first.text != '\n') {
-                this.result += response.choices.first.text;
-              }
-            });
-          },
-          onEnd: () async {
-            setState(() {
-              if (i != lines.length - 1) {
-                this.result += '\n';
-              }
-            });
-          },
-          onError: (error) async {
-            await _showError(error.message);
-            await _onExit();
-          },
-        );
-      }
-    }, (r) async {
-      await _showError(r.msg);
-      await _onExit();
-    });
+    var lines = input.split('\n\n');
+    if (action == SmartEditAction.summarize) {
+      lines = [lines.join('\n')];
+    }
+    for (var i = 0; i < lines.length; i++) {
+      final element = lines[i];
+      await openAIRepository.getStreamedCompletions(
+        useAction: true,
+        prompt: action.prompt(element),
+        onStart: () async {
+          setState(() {
+            loading = false;
+          });
+        },
+        onProcess: (response) async {
+          setState(() {
+            if (response.choices.first.text != '\n') {
+              result += response.choices.first.text;
+            }
+          });
+        },
+        onEnd: () async {
+          setState(() {
+            if (i != lines.length - 1) {
+              result += '\n';
+            }
+          });
+        },
+        onError: (error) async {
+          await _showError(error.message);
+          await _onExit();
+        },
+      );
+    }
   }
 
   Future<void> _showError(String message) async {
diff --git a/frontend/appflowy_flutter/lib/startup/deps_resolver.dart b/frontend/appflowy_flutter/lib/startup/deps_resolver.dart
index 17ae094ca4..713eb7f27c 100644
--- a/frontend/appflowy_flutter/lib/startup/deps_resolver.dart
+++ b/frontend/appflowy_flutter/lib/startup/deps_resolver.dart
@@ -4,6 +4,7 @@ import 'package:appflowy/plugins/database_view/application/field/field_controlle
 import 'package:appflowy/plugins/database_view/application/field/field_service.dart';
 import 'package:appflowy/plugins/database_view/application/setting/property_bloc.dart';
 import 'package:appflowy/plugins/database_view/grid/application/grid_header_bloc.dart';
+import 'package:appflowy/plugins/document/presentation/plugins/openai/service/openai_client.dart';
 import 'package:appflowy/user/application/user_listener.dart';
 import 'package:appflowy/user/application/user_service.dart';
 import 'package:appflowy/util/file_picker/file_picker_impl.dart';
@@ -27,6 +28,7 @@ import 'package:appflowy_backend/protobuf/flowy-folder/view.pb.dart';
 import 'package:appflowy_backend/protobuf/flowy-user/user_profile.pb.dart';
 import 'package:fluttertoast/fluttertoast.dart';
 import 'package:get_it/get_it.dart';
+import 'package:http/http.dart' as http;
 
 class DependencyResolver {
   static Future<void> resolve(GetIt getIt) async {
@@ -44,8 +46,25 @@ class DependencyResolver {
   }
 }
 
-void _resolveCommonService(GetIt getIt) {
+void _resolveCommonService(GetIt getIt) async {
   getIt.registerFactory<FilePickerService>(() => FilePicker());
+
+  getIt.registerFactoryAsync<OpenAIRepository>(
+    () async {
+      final result = await UserBackendService.getCurrentUserProfile();
+      return result.fold(
+        (l) {
+          return HttpOpenAIRepository(
+            client: http.Client(),
+            apiKey: l.openaiKey,
+          );
+        },
+        (r) {
+          throw Exception('Failed to get user profile: ${r.msg}');
+        },
+      );
+    },
+  );
 }
 
 void _resolveUserDeps(GetIt getIt) {
@@ -160,4 +179,4 @@ void _resolveGridDeps(GetIt getIt) {
     (viewId, cache) =>
         DatabasePropertyBloc(viewId: viewId, fieldController: cache),
   );
-}
+}
\ No newline at end of file
diff --git a/frontend/appflowy_flutter/pubspec.lock b/frontend/appflowy_flutter/pubspec.lock
index 1748789b39..aa31fb15fe 100644
--- a/frontend/appflowy_flutter/pubspec.lock
+++ b/frontend/appflowy_flutter/pubspec.lock
@@ -809,7 +809,7 @@ packages:
     source: hosted
     version: "1.0.4"
   mocktail:
-    dependency: transitive
+    dependency: "direct main"
     description:
       name: mocktail
       sha256: "80a996cd9a69284b3dc521ce185ffe9150cde69767c2d3a0720147d93c0cef53"
diff --git a/frontend/appflowy_flutter/pubspec.yaml b/frontend/appflowy_flutter/pubspec.yaml
index d503b8475c..f12622f15e 100644
--- a/frontend/appflowy_flutter/pubspec.yaml
+++ b/frontend/appflowy_flutter/pubspec.yaml
@@ -42,7 +42,7 @@ dependencies:
     git:
       url: https://github.com/AppFlowy-IO/appflowy-board.git
       ref: a183c57
-  appflowy_editor: "^0.1.9"
+  appflowy_editor: ^0.1.9
   appflowy_popover:
     path: packages/appflowy_popover
 
@@ -98,6 +98,7 @@ dependencies:
   http: ^0.13.5
   json_annotation: ^4.7.0
   path: ^1.8.2
+  mocktail: ^0.3.0
   archive: ^3.3.0
   flutter_svg: ^2.0.5