diff --git a/app/controllers/submissions_controller.rb b/app/controllers/submissions_controller.rb
index 1301d188..3465bfa0 100644
--- a/app/controllers/submissions_controller.rb
+++ b/app/controllers/submissions_controller.rb
@@ -178,7 +178,7 @@ class SubmissionsController < ApplicationController
         send_and_store client_socket, message
       end
 
-      runner_socket.on :exit do |exit_code, files|
+      runner_socket.on :exit do |exit_code|
         @testrun[:exit_code] = exit_code
         exit_statement =
           if @testrun[:output].empty? && exit_code.zero?
@@ -201,6 +201,10 @@ class SubmissionsController < ApplicationController
           @testrun[:status] = :out_of_memory
         end
 
+        # The client connection will be closed once the file listing finished.
+      end
+
+      runner_socket.on :files do |files|
         downloadable_files, = convert_files_json_to_files files
         if downloadable_files.present?
           js_tree = FileTree.new(downloadable_files).to_js_tree
diff --git a/lib/runner/connection.rb b/lib/runner/connection.rb
index ad23ae5b..df8f687f 100644
--- a/lib/runner/connection.rb
+++ b/lib/runner/connection.rb
@@ -5,7 +5,7 @@ require 'json_schemer'
 
 class Runner::Connection
   # These are events for which callbacks can be registered.
-  EVENTS = %i[start exit stdout stderr].freeze
+  EVENTS = %i[start exit stdout stderr files].freeze
   WEBSOCKET_MESSAGE_TYPES = %i[start stdout stderr error timeout exit].freeze
   BACKEND_OUTPUT_SCHEMA = JSONSchemer.schema(JSON.parse(File.read('lib/runner/backend-output.schema.json')))
   SENTRY_OP_NAME = 'websocket.client'
@@ -15,6 +15,7 @@ class Runner::Connection
   # @!attribute exit_callback
   # @!attribute stdout_callback
   # @!attribute stderr_callback
+  # @!attribute files_callback
   attr_writer :status
   attr_reader :error
 
@@ -158,12 +159,8 @@ class Runner::Connection
         @strategy.destroy_at_management
         @error = Runner::Error::ExecutionTimeout.new('Execution exceeded its time limit')
       when :terminated_by_codeocean, :terminated_by_management
-        files = begin
-          @strategy.retrieve_files
-        rescue Runner::Error::RunnerNotFound, Runner::Error::WorkspaceError
-          {'files' => []}
-        end
-        @exit_callback.call @exit_code, files
+        @exit_callback.call @exit_code
+        list_filesystem
       when :terminated_by_client, :error
         @strategy.destroy_at_management
       else # :established
@@ -181,6 +178,20 @@ class Runner::Connection
     @event_loop.stop
   end
 
+  def list_filesystem
+    files = {'files' => []}
+    begin
+      # Retrieve files from runner management ONLY IF the callback was defined outside of this class.
+      # Otherwise, we would call our default callback and retrieve the files without any further processing.
+      files = @strategy.retrieve_files if @files_callback.source_location.first != __FILE__
+    rescue Runner::Error::RunnerNotFound, Runner::Error::WorkspaceError
+      # Ignore errors when retrieving files. This is not critical and a suitable default is already provided.
+    ensure
+      @files_callback.call files
+    end
+  end
+  private :list_filesystem
+
   # === Message Handlers
   # Each message type indicated by the +type+ attribute in the JSON
   # sent be the runner management has a dedicated method.