From 94b2f28dd74c1d19b6048892936deda26a6f239e Mon Sep 17 00:00:00 2001
From: DerThorsten <derthorstenbeier@gmail.com>
Date: Wed, 20 Mar 2024 10:24:19 +0100
Subject: [PATCH] handle injection of execution_count

---
 include/xeus/xrequest_context.hpp |  3 +++
 src/xinterpreter.cpp              | 12 ++++++++++--
 src/xrequest_context.cpp          | 12 ++++++++++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/include/xeus/xrequest_context.hpp b/include/xeus/xrequest_context.hpp
index bf4d6485..d797808a 100644
--- a/include/xeus/xrequest_context.hpp
+++ b/include/xeus/xrequest_context.hpp
@@ -48,6 +48,9 @@ namespace xeus
         xexecute_request_context(nl::json header, channel origin, guid_list id, std::function<void(const xexecute_request_context&,nl::json)> on_send_reply);
 
         void send_reply(nl::json reply);
+        void set_on_send_callback(std::function<void(const xexecute_request_context&,nl::json)> augment_reply);
+        std::function<void(const xexecute_request_context&, nl::json)> get_on_send_callback() const;
+
         private:
         std::function<void(const xexecute_request_context&, nl::json)> m_on_send_reply;
     };
diff --git a/src/xinterpreter.cpp b/src/xinterpreter.cpp
index 0263707d..40a6bbf9 100644
--- a/src/xinterpreter.cpp
+++ b/src/xinterpreter.cpp
@@ -41,6 +41,16 @@ namespace xeus
             ++m_execution_count;
             publish_execution_input(context, code, m_execution_count);
         }
+        // copy m_execution_count in a local variable to capture it in the lambda
+        auto execution_count = m_execution_count;
+
+        auto inner_callback = context.get_on_send_callback();
+        context.set_on_send_callback( [inner_callback, execution_count](const xexecute_request_context& context,nl::json reply)
+        {
+            reply["execution_count"] = execution_count;
+            inner_callback(context, reply);
+        });
+
 
         execute_request_impl(
             std::move(context),
@@ -48,8 +58,6 @@ namespace xeus
             store_history, user_expressions, allow_stdin
         );
 
-        // reply["execution_count"] = m_execution_count;
-        // return reply;
     }
 
     nl::json xinterpreter::complete_request(const std::string& code, int cursor_pos)
diff --git a/src/xrequest_context.cpp b/src/xrequest_context.cpp
index 6fd38295..917122c6 100644
--- a/src/xrequest_context.cpp
+++ b/src/xrequest_context.cpp
@@ -39,4 +39,16 @@ namespace xeus
           m_on_send_reply(std::move(on_send_reply))
     {
     }
+
+    // get and set callback
+    void xexecute_request_context::set_on_send_callback(std::function<void(const xexecute_request_context&,nl::json)> augment_reply)
+    {
+        m_on_send_reply = std::move(augment_reply);
+    }
+
+    // get callback
+    std::function<void(const xexecute_request_context&, nl::json)> xexecute_request_context::get_on_send_callback() const
+    {
+        return m_on_send_reply;
+    }
 }
\ No newline at end of file