diff --git a/MIGRATION.md b/MIGRATION.md
index 7ae86b8f..63a87e19 100644
--- a/MIGRATION.md
+++ b/MIGRATION.md
@@ -17,7 +17,8 @@ Also make sure to run `:UpdateRemotePlugins` to cleanup the old python commands.
 Removed or changed params that you pass to `setup`:
 
 - `show_help` was removed (the help is now always shown as virtual text, and not intrusive)
-- `hide_system_prompt` was removed (it is part of `disable_extra_info`)
+- `disable_extra_info` was renamed to `show_user_selection`
+- `hide_system_prompt` was renamed to `show_system_prompt`
 - `proxy` does not work at the moment (waiting for change in plenary.nvim), if you are behind corporate proxy you can look at something like [vpn-slice](https://github.com/dlenski/vpn-slice)
 - `language` was removed and is now part of `selection` as `selection.filetype`
 
@@ -58,8 +59,8 @@ end, { nargs = '*', range = true })
 - [ ] For proxy support, this is needed: https://github.com/nvim-lua/plenary.nvim/pull/559
 - [ ] Delete rest of the python code? Or finish rewriting in place then delete
 - [ ] Check for curl availability with health check
-- [ ] Add folds logic from python, maybe? Not sure if this is even needed
+- [x] Add folds logic from python, maybe? Not sure if this is even needed
 - [ ] As said in changes part, finish rewriting the authentication request if needed
-- [ ] Properly get token file path, atm it only supports Linux (easy fix)
+- [x] Properly get token file path, atm it only supports Linux (easy fix)
 - [ ] Update README and stuff
 - [ ] Add token count support to extra_info, something like this called from lua:
diff --git a/lua/CopilotChat/copilot.lua b/lua/CopilotChat/copilot.lua
index 5971721a..4fedf141 100644
--- a/lua/CopilotChat/copilot.lua
+++ b/lua/CopilotChat/copilot.lua
@@ -122,9 +122,8 @@ local function authenticate(github_token)
   return sessionid, token, nil
 end
 
-local Copilot = class(function(self, show_extra_info)
+local Copilot = class(function(self)
   self.github_token = get_cached_token()
-  self.show_extra_info = show_extra_info or false
   self.history = {}
   self.token = nil
   self.sessionid = nil
@@ -182,16 +181,6 @@ function Copilot:ask(prompt, opts)
   log.debug('Model: ' .. model)
   log.debug('Temperature: ' .. temperature)
 
-  if self.show_extra_info and on_progress then
-    on_progress('SYSTEM PROMPT:\n```\n' .. system_prompt .. '```\n')
-
-    if selection ~= '' then
-      on_progress('CODE:\n```' .. filetype .. '\n' .. selection .. '\n```')
-    end
-
-    on_done('')
-  end
-
   table.insert(self.history, {
     content = prompt,
     role = 'user',
@@ -202,14 +191,6 @@ function Copilot:ask(prompt, opts)
     self:stop()
   end
 
-  -- Notify the user about current prompt
-  if on_progress then
-    on_progress(prompt)
-  end
-  if on_done then
-    on_done(prompt)
-  end
-
   if on_start then
     on_start()
   end
@@ -262,7 +243,9 @@ function Copilot:ask(prompt, opts)
 
         if not ok then
           log.error('Failed parse response: ' .. tostring(err))
-          on_error(content)
+          if on_error then
+            on_error(content)
+          end
           return
         end
 
diff --git a/lua/CopilotChat/init.lua b/lua/CopilotChat/init.lua
index fa9f6ec7..8268a662 100644
--- a/lua/CopilotChat/init.lua
+++ b/lua/CopilotChat/init.lua
@@ -6,6 +6,7 @@ local select = require('CopilotChat.select')
 local debuginfo = require('CopilotChat.debuginfo')
 
 local M = {}
+local plugin_name = 'CopilotChat.nvim'
 local state = {
   copilot = nil,
   chat = nil,
@@ -13,6 +14,16 @@ local state = {
   window = nil,
 }
 
+function CopilotChatFoldExpr(lnum, separator)
+  local line = vim.fn.getline(lnum)
+  vim.print(line)
+  if string.match(line, separator .. '$') then
+    return '>1'
+  end
+
+  return '='
+end
+
 local function find_lines_between_separator_at_cursor(bufnr, separator)
   local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
   local cursor = vim.api.nvim_win_get_cursor(0)
@@ -175,7 +186,7 @@ function M.open(config)
   local just_created = false
 
   if not state.chat or not state.chat:valid() then
-    state.chat = Chat(M.config.name)
+    state.chat = Chat(plugin_name)
     just_created = true
 
     if config.mappings.complete then
@@ -193,8 +204,8 @@ function M.open(config)
     if config.mappings.submit_prompt then
       vim.keymap.set('n', config.mappings.submit_prompt, function()
         local input, start_line, end_line, line_count =
-          find_lines_between_separator_at_cursor(state.chat.bufnr, M.config.separator)
-        if input ~= '' and not vim.startswith(vim.trim(input), '**' .. M.config.name .. ':**') then
+          find_lines_between_separator_at_cursor(state.chat.bufnr, config.separator)
+        if input ~= '' and not vim.startswith(vim.trim(input), '**' .. config.name .. ':**') then
           -- If we are entering the input at the end, replace it
           if line_count == end_line then
             vim.api.nvim_buf_set_lines(state.chat.bufnr, start_line, end_line, false, { '' })
@@ -265,6 +276,15 @@ function M.open(config)
     vim.wo[state.window].cursorline = true
     vim.wo[state.window].conceallevel = 2
     vim.wo[state.window].concealcursor = 'niv'
+    if config.show_folds then
+      vim.wo[state.window].foldcolumn = '1'
+      vim.wo[state.window].foldmethod = 'expr'
+      vim.wo[state.window].foldexpr = "v:lua.CopilotChatFoldExpr(v:lnum, '"
+        .. config.separator
+        .. "')"
+    else
+      vim.wo[state.window].foldcolumn = '0'
+    end
 
     if just_created then
       M.reset()
@@ -327,7 +347,26 @@ function M.ask(prompt, config)
     updated_prompt = updated_prompt .. ' ' .. state.selection.prompt_extra
   end
 
-  local just_started = true
+  local finish = false
+  if config.show_system_prompt then
+    finish = true
+    append(' **System prompt** ---\n```\n' .. system_prompt .. '```\n')
+  end
+  if config.show_user_selection and state.selection and state.selection.lines ~= '' then
+    finish = true
+    append(
+      ' **Selection** ---\n```'
+        .. (state.selection.filetype or '')
+        .. '\n'
+        .. state.selection.lines
+        .. '\n```'
+    )
+  end
+  if finish then
+    append('\n' .. config.separator .. '\n\n')
+  end
+
+  append(updated_prompt)
 
   return state.copilot:ask(updated_prompt, {
     selection = state.selection.lines,
@@ -336,21 +375,14 @@ function M.ask(prompt, config)
     model = config.model,
     temperature = config.temperature,
     on_start = function()
+      append('\n\n **' .. config.name .. '** ' .. config.separator .. '\n\n')
       state.chat.spinner:start()
-      append('**' .. M.config.name .. ':** ')
-      just_started = true
     end,
     on_done = function()
-      append('\n\n' .. M.config.separator .. '\n\n')
+      append('\n\n' .. config.separator .. '\n\n')
       show_help()
     end,
     on_progress = function(token)
-      -- Add a newline if the token is not a word and we just started (for example code block)
-      if just_started and not token:match('^%w') then
-        token = '\n' .. token
-      end
-
-      just_started = false
       append(token)
     end,
   })
@@ -370,9 +402,11 @@ M.config = {
   system_prompt = prompts.COPILOT_INSTRUCTIONS,
   model = 'gpt-4',
   temperature = 0.1,
-  debug = false,
-  clear_chat_on_new_prompt = false,
-  disable_extra_info = true,
+  debug = false, -- Enable debug logging
+  show_user_selection = true, -- Shows user selection in chat
+  show_system_prompt = false, -- Shows system prompt in chat
+  show_folds = true, -- Shows folds for sections in chat
+  clear_chat_on_new_prompt = false, -- Clears chat on every new prompt
   name = 'CopilotChat',
   separator = '---',
   prompts = {
@@ -422,12 +456,12 @@ M.config = {
 --       - mappings: (table?).
 function M.setup(config)
   M.config = vim.tbl_deep_extend('force', M.config, config or {})
-  state.copilot = Copilot(not M.config.disable_extra_info)
+  state.copilot = Copilot()
   debuginfo.setup()
 
-  local logfile = string.format('%s/%s.log', vim.fn.stdpath('state'), 'CopilotChat.nvim')
+  local logfile = string.format('%s/%s.log', vim.fn.stdpath('state'), plugin_name)
   log.new({
-    plugin = M.config.name,
+    plugin = plugin_name,
     level = M.config.debug and 'trace' or 'warn',
     outfile = logfile,
   }, true)
@@ -444,13 +478,13 @@ function M.setup(config)
       nargs = '*',
       force = true,
       range = true,
-      desc = prompt.description or ('CopilotChat.nvim ' .. name),
+      desc = prompt.description or (plugin_name .. ' ' .. name),
     })
 
     if prompt.mapping then
       vim.keymap.set({ 'n', 'v' }, prompt.mapping, function()
         M.ask(prompt.prompt, prompt)
-      end, { desc = prompt.description or ('CopilotChat.nvim ' .. name) })
+      end, { desc = prompt.description or (plugin_name .. ' ' .. name) })
     end
   end