correct subtle multi threading issues

Author: SamSaffron

app/controllers/discourse_ai/admin/ai_personas_controller.rb

@@ -90,7 +90,11 @@ def queue_reply(&block)
 
           db = RailsMultisite::ConnectionManagement.current_db
           thread_pool.post do
-            RailsMultisite::ConnectionManagement.with_connection(db) { block.call }
+            begin
+              RailsMultisite::ConnectionManagement.with_connection(db) { block.call }
+            rescue StandardError => e
+              Discourse.warn_exception(e, message: "Discourse AI: Unable to stream reply")
+            end
           end
         end
       end
@@ -157,7 +161,7 @@ def stream_reply
             PostCreator.create!(
               user,
               title: I18n.t("discourse_ai.ai_bot.default_pm_prefix"),
-              raw: params[:question],
+              raw: params[:query],
               archetype: Archetype.private_message,
               target_usernames: "#{user.username},#{persona.user.username}",
               skip_validations: true,
@@ -175,16 +179,18 @@ def stream_reply
           begin
             io.write "HTTP/1.1 200 OK"
             io.write CRLF
-            io.write "Content-Type: application/json"
+            io.write "Content-Type: text/plain; charset=utf-8"
             io.write CRLF
             io.write "Transfer-Encoding: chunked"
             io.write CRLF
             io.write "Cache-Control: no-cache, no-store, must-revalidate"
             io.write CRLF
-            io.write "Connection: keep-alive"
+            io.write "Connection: close"
             io.write CRLF
             io.write "X-Accel-Buffering: no"
             io.write CRLF
+            io.write "X-Content-Type-Options: nosniff"
+            io.write CRLF
             io.write CRLF
             io.flush
 
@@ -202,31 +208,33 @@ def stream_reply
             io.write data
             io.write CRLF
 
-            io.flush
-
             DiscourseAi::AiBot::Playground
               .new(bot)
               .reply_to(post) do |partial|
+                next if partial.length == 0
+
                 data = { partial: partial }.to_json + "\n\n"
 
+                data.force_encoding("UTF-8")
+
                 io.write data.bytesize.to_s(16)
                 io.write CRLF
                 io.write data
                 io.write CRLF
-
                 io.flush
               end
 
-            # End the response with zero-length chunk
             io.write "0"
             io.write CRLF
             io.write CRLF
+
             io.flush
+            io.done
           rescue StandardError => e
-            p e
-            puts e.backtrace
-          rescue IOError => e
-            Rails.logger.error "Streaming error: #{e.message}"
+            # make it a tiny bit easier to debug in dev, this is tricky
+            # multi-threaded code that exhibits various limitations in rails
+            p e if Rails.env.development?
+            Discourse.warn_exception(e, message: "Discourse AI: Unable to stream reply")
           ensure
             io.close
           end

config/locales/server.en.yml

@@ -169,6 +169,7 @@ en:
         failed_to_share: "Failed to share the conversation"
         conversation_deleted: "Conversation share deleted successfully"
     ai_bot:
+      default_pm_prefix: "[Untitled AI bot PM]"
       personas:
         default_llm_required: "Default LLM model is required prior to enabling Chat"
         cannot_delete_system_persona: "System personas cannot be deleted, please disable it instead"

lib/ai_bot/playground.rb

@@ -391,6 +391,11 @@ def get_context(participants:, conversation_context:, user:, skip_tool_details:
       end
 
       def reply_to(post, &blk)
+        # this is a multithreading issue
+        # post custom prompt is needed and it may not
+        # be properly loaded, ensure it is loaded
+        PostCustomPrompt.none
+
         reply = +""
         start = Time.now
 

spec/requests/admin/ai_personas_controller_spec.rb

@@ -456,11 +456,12 @@ def validate_streamed_response(raw_http, expected)
 
       preamble = (<<~PREAMBLE).strip
         HTTP/1.1 200 OK
-        Content-Type: application/json
+        Content-Type: text/plain; charset=utf-8
         Transfer-Encoding: chunked
         Cache-Control: no-cache, no-store, must-revalidate
-        Connection: keep-alive
+        Connection: close
         X-Accel-Buffering: no
+        X-Content-Type-Options: nosniff
       PREAMBLE
 
       expect(header_lines.join("\n")).to eq(preamble)
@@ -522,6 +523,9 @@ def validate_streamed_response(raw_http, expected)
       last_post = topic.posts.order(:created_at).last
       expect(last_post.raw).to eq("This is a test! Testing!")
 
+      user_post = topic.posts.find_by(post_number: 1)
+      expect(user_post.raw).to eq("how are you today?")
+
       # need ai persona and user
       expect(topic.topic_allowed_users.count).to eq(2)
       expect(topic.archetype).to eq(Archetype.private_message)