diff --git a/lua/cp/setup.lua b/lua/cp/setup.lua
index 6dd579f..12fde62 100644
--- a/lua/cp/setup.lua
+++ b/lua/cp/setup.lua
@@ -172,6 +172,7 @@ function M.setup_problem(problem_id, language)
       vim.api.nvim_buf_set_name(prov.bufnr, source_file)
       vim.bo[prov.bufnr].swapfile = true
       vim.cmd(string.format('silent keepalt noautocmd write! %s', vim.fn.fnameescape(source_file)))
+      state.set_solution_win(vim.api.nvim_get_current_win())
       if config.hooks and config.hooks.setup_code and not vim.b[prov.bufnr].cp_setup_done then
         local ok = pcall(config.hooks.setup_code, state)
         if ok then
@@ -194,6 +195,7 @@ function M.setup_problem(problem_id, language)
     vim.cmd.only({ mods = { silent = true } })
     vim.cmd.e(source_file)
     local bufnr = vim.api.nvim_get_current_buf()
+    state.set_solution_win(vim.api.nvim_get_current_win())
     if config.hooks and config.hooks.setup_code and not vim.b[bufnr].cp_setup_done then
       local ok = pcall(config.hooks.setup_code, state)
       if ok then
diff --git a/lua/cp/state.lua b/lua/cp/state.lua
index 63613a3..c234387 100644
--- a/lua/cp/state.lua
+++ b/lua/cp/state.lua
@@ -35,6 +35,7 @@ local state = {
   saved_session = nil,
   active_panel = nil,
   provisional = nil,
+  solution_win = nil,
 }
 
 ---@return string|nil
@@ -153,6 +154,19 @@ function M.set_provisional(p)
   state.provisional = p
 end
 
+---@return integer?
+function M.get_solution_win()
+  if state.solution_win and vim.api.nvim_win_is_valid(state.solution_win) then
+    return state.solution_win
+  end
+  return vim.api.nvim_get_current_win()
+end
+
+---@param win integer?
+function M.set_solution_win(win)
+  state.solution_win = win
+end
+
 M._state = state
 
 return M