new cancel manager makes cancelling request much easier

Author: SamSaffron

lib/ai_bot/playground.rb

@@ -471,6 +471,13 @@ def reply_to(
 
           redis_stream_key = "gpt_cancel:#{reply_post.id}"
           Discourse.redis.setex(redis_stream_key, MAX_STREAM_DELAY_SECONDS, 1)
+
+          context.cancel_manager = DiscourseAi::Completions::CancelManager.new
+          context
+            .cancel_manager
+            .start_monitor(delay: 0.2) do
+              context.cancel_manager.cancel! if !Discourse.redis.get(redis_stream_key)
+            end
         end
 
         context.skip_tool_details ||= !bot.persona.class.tool_details
@@ -568,6 +575,8 @@ def reply_to(
         end
         raise e
       ensure
+        context.cancel_manager.stop_monitor if context&.cancel_manager
+
         # since we are skipping validations and jobs we
         # may need to fix participant count
         if reply_post && reply_post.topic && reply_post.topic.private_message? &&
@@ -649,7 +658,7 @@ def publish_update(bot_reply_post, payload)
           payload,
           user_ids: bot_reply_post.topic.allowed_user_ids,
           max_backlog_size: 2,
-          max_backlog_age: 60,
+          max_backlog_age: MAX_STREAM_DELAY_SECONDS,
         )
       end
     end

lib/completions/cancel_manager.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+# special object that can be used to cancel completions and http requests
+module DiscourseAi
+  module Completions
+    class CancelManager
+      attr_reader :cancelled
+      attr_reader :callbacks
+
+      def initialize
+        @cancelled = false
+        @callbacks = Concurrent::Array.new
+        @mutex = Mutex.new
+        @monitor_thread = nil
+      end
+
+      def monitor_thread
+        @mutex.synchronize { @monitor_thread }
+      end
+
+      def start_monitor(delay: 0.5, &block)
+        @mutex.synchronize do
+          raise "Already monitoring" if @monitor_thread
+          raise "Expected a block" if !block
+
+          db = RailsMultisite::ConnectionManagement.current_db
+          @stop_monitor = false
+
+          @monitor_thread =
+            Thread.new do
+              begin
+                loop do
+                  @mutex.synchronize { break if @stop_monitor }
+                  sleep delay
+                  @mutex.synchronize { break if @stop_monitor }
+                  @mutex.synchronize { break if cancelled? }
+
+                  should_cancel = false
+                  RailsMultisite::ConnectionManagement.with_connection(db) do
+                    should_cancel = block.call
+                  end
+
+                  @mutex.synchronize { cancel! if should_cancel }
+
+                  break if cancelled?
+                end
+              ensure
+                @mutex.synchronize { @monitor_thread = nil }
+              end
+            end
+        end
+      end
+
+      def stop_monitor
+        monitor_thread = nil
+
+        @mutex.synchronize { monitor_thread = @monitor_thread }
+
+        if monitor_thread
+          @mutex.synchronize { @stop_monitor = true }
+          # so we do not deadlock
+          monitor_thread.wakeup
+          monitor_thread.join(2)
+          # should not happen
+          if monitor_thread.alive?
+            Rails.logger.warn("DiscourseAI: CancelManager monitor thread did not stop in time")
+            monitor_thread.kill if monitor_thread.alive?
+          end
+        end
+      end
+
+      def cancelled?
+        @cancelled
+      end
+
+      def add_callback(cb)
+        @callbacks << cb
+      end
+
+      def remove_callback(cb)
+        @callbacks.delete(cb)
+      end
+
+      def cancel!
+        @cancelled = true
+        @callbacks.each do |cb|
+          begin
+            cb.call
+          rescue StandardError
+            # ignore cause this may have already been cancelled
+          end
+        end
+      end
+    end
+  end
+end

lib/completions/endpoints/base.rb

@@ -68,11 +68,17 @@ def perform_completion!(
           feature_context: nil,
           partial_tool_calls: false,
           output_thinking: false,
+          cancel_manager: nil,
           &blk
         )
           LlmQuota.check_quotas!(@llm_model, user)
           start_time = Time.now
 
+          if cancel_manager && cancel_manager.cancelled?
+            # nothing to do
+            return
+          end
+
           @forced_json_through_prefill = false
           @partial_tool_calls = partial_tool_calls
           @output_thinking = output_thinking
@@ -90,6 +96,7 @@ def perform_completion!(
                 feature_context: feature_context,
                 partial_tool_calls: partial_tool_calls,
                 output_thinking: output_thinking,
+                cancel_manager: cancel_manager,
               )
 
             wrapped = result
@@ -118,6 +125,9 @@ def perform_completion!(
             end
           end
 
+          cancel_manager_callback = nil
+          cancelled = false
+
           FinalDestination::HTTP.start(
             model_uri.host,
             model_uri.port,
@@ -126,6 +136,14 @@ def perform_completion!(
             open_timeout: TIMEOUT,
             write_timeout: TIMEOUT,
           ) do |http|
+            if cancel_manager
+              cancel_manager_callback =
+                lambda do
+                  cancelled = true
+                  http.finish
+                end
+              cancel_manager.add_callback(cancel_manager_callback)
+            end
             response_data = +""
             response_raw = +""
 
@@ -196,7 +214,6 @@ def perform_completion!(
               end
 
               begin
-                cancelled = false
                 cancel = -> do
                   cancelled = true
                   http.finish
@@ -224,8 +241,6 @@ def perform_completion!(
                     partials.each { |inner_partial| blk.call(inner_partial, cancel) }
                   end
                 end
-              rescue IOError, StandardError
-                raise if !cancelled
               end
               if xml_stripper
                 stripped = xml_stripper.finish
@@ -293,6 +308,12 @@ def perform_completion!(
               end
             end
           end
+        rescue IOError, StandardError
+          raise if !cancelled
+        ensure
+          if cancel_manager && cancel_manager_callback
+            cancel_manager.remove_callback(cancel_manager_callback)
+          end
         end
 
         def final_log_update(log)

lib/completions/endpoints/canned_response.rb

@@ -30,7 +30,8 @@ def perform_completion!(
           feature_name: nil,
           feature_context: nil,
           partial_tool_calls: false,
-          output_thinking: false
+          output_thinking: false,
+          cancel_manager: nil
         )
           @dialect = dialect
           @model_params = model_params

lib/completions/endpoints/fake.rb

@@ -122,7 +122,8 @@ def perform_completion!(
           feature_name: nil,
           feature_context: nil,
           partial_tool_calls: false,
-          output_thinking: false
+          output_thinking: false,
+          cancel_manager: nil
         )
           last_call = { dialect: dialect, user: user, model_params: model_params }
           self.class.last_call = last_call

lib/completions/endpoints/open_ai.rb

@@ -46,6 +46,7 @@ def perform_completion!(
           feature_context: nil,
           partial_tool_calls: false,
           output_thinking: false,
+          cancel_manager: nil,
           &blk
         )
           @disable_native_tools = dialect.disable_native_tools?

lib/completions/llm.rb

@@ -325,6 +325,7 @@ def generate(
         output_thinking: false,
         response_format: nil,
         extra_model_params: nil,
+        cancel_manager: nil,
         &partial_read_blk
       )
         self.class.record_prompt(
@@ -378,6 +379,7 @@ def generate(
           feature_context: feature_context,
           partial_tool_calls: partial_tool_calls,
           output_thinking: output_thinking,
+          cancel_manager: cancel_manager,
           &partial_read_blk
         )
       end

lib/personas/bot.rb

@@ -55,6 +55,7 @@ def reply(context, llm_args: {}, &update_blk)
         unless context.is_a?(BotContext)
           raise ArgumentError, "context must be an instance of BotContext"
         end
+        context.cancel_manager ||= DiscourseAi::Completions::CancelManager.new
         current_llm = llm
         prompt = persona.craft_prompt(context, llm: current_llm)
 
@@ -91,6 +92,7 @@ def reply(context, llm_args: {}, &update_blk)
               feature_name: context.feature_name,
               partial_tool_calls: allow_partial_tool_calls,
               output_thinking: true,
+              cancel_manager: context.cancel_manager,
               **llm_kwargs,
             ) do |partial, cancel|
               tool =

lib/personas/bot_context.rb

@@ -16,7 +16,8 @@ class BotContext
                     :channel_id,
                     :context_post_ids,
                     :feature_name,
-                    :resource_url
+                    :resource_url,
+                    :cancel_manager
 
       def initialize(
         post: nil,

lib/personas/tools/researcher.rb

@@ -101,9 +101,12 @@ def run_inference(chunk_text, goal, post, &blk)
             )
 
           results = []
-          llm.generate(prompt, user: post.user, feature_name: context.feature_name) do |partial|
-            results << partial
-          end
+          llm.generate(
+            prompt,
+            user: post.user,
+            feature_name: context.feature_name,
+            cancel_manager: context.cancel_manager,
+          ) { |partial| results << partial }
 
           blk.call(".")
           results.join