Treat 'exec' & 'include' roles same as '>>> user'

Author: juodumas

py/chat.py

@@ -10,8 +10,8 @@ def run_ai_chat(context):
     config_ui = config['ui']
 
     def initialize_chat_window():
-        lines = vim.eval('getline(1, "$")')
-        contains_user_prompt = '>>> user' in lines
+        file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
+        contains_user_prompt = re.search(r"^>>> (user|exec|include)", file_content, flags=re.MULTILINE)
         if not contains_user_prompt:
             # user role not found, put whole file content as an user prompt
             vim.command("normal! gg")
@@ -30,10 +30,9 @@ def initialize_chat_window():
         vim_break_undo_sequence()
         vim.command("redraw")
 
-        file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
-        role_lines = re.findall(r'(^>>> user|^>>> system|^<<< thinking|^<<< assistant).*', file_content, flags=re.MULTILINE)
-        if not role_lines[-1].startswith(">>> user"):
-            # last role is not user, most likely completion was cancelled before
+        last_role = re.match(r".*^>>> (\w+)", file_content, flags=re.DOTALL | re.MULTILINE)
+        if last_role and last_role.group(1) not in ('user', 'include', 'exec'):
+            # last role is not a user role, most likely completion was cancelled before
             vim.command("normal! o")
             vim.command("normal! i\n>>> user\n\n")