diff --git a/lua/cp/window.lua b/lua/cp/window.lua
index 1641f47..b150965 100644
--- a/lua/cp/window.lua
+++ b/lua/cp/window.lua
@@ -75,10 +75,11 @@ function M.restore_layout(state, tile_fn)
     local source_file
     if source_files ~= '' then
       local files = vim.split(source_files, '\n')
-      local valid_extensions = vim.tbl_keys(constants.filetype_to_language)
+      -- Prefer known extensions first, but accept any extension
+      local known_extensions = vim.tbl_keys(constants.filetype_to_language)
       for _, file in ipairs(files) do
         local ext = vim.fn.fnamemodify(file, ':e')
-        if vim.tbl_contains(valid_extensions, ext) then
+        if vim.tbl_contains(known_extensions, ext) then
           source_file = file
           break
         end
diff --git a/spec/command_parsing_spec.lua b/spec/command_parsing_spec.lua
index ebdb07c..15a1cd2 100644
--- a/spec/command_parsing_spec.lua
+++ b/spec/command_parsing_spec.lua
@@ -32,7 +32,7 @@ describe('cp command parsing', function()
   end)
 
   describe('empty arguments', function()
-    it('logs error for no arguments', function()
+    it('attempts file state restoration for no arguments', function()
       local opts = { fargs = {} }
 
       cp.handle_command(opts)
@@ -40,7 +40,10 @@ describe('cp command parsing', function()
       assert.is_true(#logged_messages > 0)
       local error_logged = false
       for _, log_entry in ipairs(logged_messages) do
-        if log_entry.level == vim.log.levels.ERROR and log_entry.msg:match('Usage:') then
+        if
+          log_entry.level == vim.log.levels.ERROR
+          and log_entry.msg:match('No file is currently open')
+        then
           error_logged = true
           break
         end