streamline cancel logic

Author: SamSaffron

lib/ai_bot/chat_streamer.rb

@@ -6,16 +6,23 @@
 module DiscourseAi
   module AiBot
     class ChatStreamer
-      attr_accessor :cancel
       attr_reader :reply,
                   :guardian,
                   :thread_id,
                   :force_thread,
                   :in_reply_to_id,
                   :channel,
-                  :cancelled
-
-      def initialize(message:, channel:, guardian:, thread_id:, in_reply_to_id:, force_thread:)
+                  :cancel_manager
+
+      def initialize(
+        message:,
+        channel:,
+        guardian:,
+        thread_id:,
+        in_reply_to_id:,
+        force_thread:,
+        cancel_manager: nil
+      )
         @message = message
         @channel = channel
         @guardian = guardian
@@ -35,6 +42,8 @@ def initialize(message:, channel:, guardian:, thread_id:, in_reply_to_id:, force
             guardian: guardian,
             thread_id: thread_id,
           )
+
+        @cancel_manager = cancel_manager
       end
 
       def <<(partial)
@@ -111,8 +120,7 @@ def run
 
           streaming = ChatSDK::Message.stream(message_id: reply.id, raw: buffer, guardian: guardian)
           if !streaming
-            cancel.call
-            @cancelled = true
+            @cancel_manager.cancel! if @cancel_manager
           end
         end
       end

lib/ai_bot/playground.rb

@@ -331,6 +331,7 @@ def reply_to_chat_message(message, channel, context_post_ids)
               ),
             user: message.user,
             skip_tool_details: true,
+            cancel_manager: DiscourseAi::Completions::CancelManager.new,
           )
 
         reply = nil
@@ -347,15 +348,14 @@ def reply_to_chat_message(message, channel, context_post_ids)
             thread_id: message.thread_id,
             in_reply_to_id: in_reply_to_id,
             force_thread: force_thread,
+            cancel_manager: context.cancel_manager,
           )
 
         new_prompts =
-          bot.reply(context) do |partial, cancel, placeholder, type|
+          bot.reply(context) do |partial, placeholder, type|
             # no support for tools or thinking by design
             next if type == :thinking || type == :tool_details || type == :partial_tool
-            streamer.cancel = cancel
             streamer << partial
-            break if streamer.cancelled
           end
 
         reply = streamer.reply
@@ -383,6 +383,7 @@ def reply_to(
         auto_set_title: true,
         silent_mode: false,
         feature_name: nil,
+        cancel_manager: nil,
         &blk
       )
         # this is a multithreading issue
@@ -472,22 +473,25 @@ 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
+          cancel_manager ||= DiscourseAi::Completions::CancelManager.new
+          context.cancel_manager = cancel_manager
           context
             .cancel_manager
             .start_monitor(delay: 0.2) do
               context.cancel_manager.cancel! if !Discourse.redis.get(redis_stream_key)
             end
+
+          context.cancel_manager.add_callback(
+            lambda { reply_post.update!(raw: reply, cooked: PrettyText.cook(reply)) },
+          )
         end
 
         context.skip_tool_details ||= !bot.persona.class.tool_details
-
         post_streamer = PostStreamer.new(delay: Rails.env.test? ? 0 : 0.5) if stream_reply
-
         started_thinking = false
 
         new_custom_prompts =
-          bot.reply(context) do |partial, cancel, placeholder, type|
+          bot.reply(context) do |partial, placeholder, type|
             if type == :thinking && !started_thinking
               reply << "<details><summary>#{I18n.t("discourse_ai.ai_bot.thinking")}</summary>"
               started_thinking = true
@@ -506,15 +510,6 @@ def reply_to(
               blk.call(partial)
             end
 
-            if stream_reply && !Discourse.redis.get(redis_stream_key)
-              cancel&.call
-              reply_post.update!(raw: reply, cooked: PrettyText.cook(reply))
-              # we do not break out, cause if we do
-              # we will not get results from bot
-              # leading to broken context
-              # we need to trust it to cancel at the endpoint
-            end
-
             if post_streamer
               post_streamer.run_later do
                 Discourse.redis.expire(redis_stream_key, MAX_STREAM_DELAY_SECONDS)

lib/completions/endpoints/base.rb

@@ -101,11 +101,9 @@ def perform_completion!(
 
             wrapped = result
             wrapped = [result] if !result.is_a?(Array)
-            cancelled_by_caller = false
-            cancel_proc = -> { cancelled_by_caller = true }
             wrapped.each do |partial|
-              blk.call(partial, cancel_proc)
-              break if cancelled_by_caller
+              blk.call(partial)
+              break cancel_manager&.cancelled?
             end
             return result
           end
@@ -176,7 +174,7 @@ def perform_completion!(
 
               if @streaming_mode
                 blk =
-                  lambda do |partial, cancel|
+                  lambda do |partial|
                     if partial.is_a?(String)
                       partial = xml_stripper << partial if xml_stripper
 
@@ -185,7 +183,7 @@ def perform_completion!(
                         partial = structured_output
                       end
                     end
-                    orig_blk.call(partial, cancel) if partial
+                    orig_blk.call(partial) if partial
                   end
               end
 
@@ -214,13 +212,6 @@ def perform_completion!(
               end
 
               begin
-                cancel = -> do
-                  cancelled = true
-                  http.finish
-                end
-
-                break if cancelled
-
                 response.read_body do |chunk|
                   break if cancelled
 
@@ -233,12 +224,9 @@ def perform_completion!(
                     partials = [partial]
                     if xml_tool_processor && partial.is_a?(String)
                       partials = (xml_tool_processor << partial)
-                      if xml_tool_processor.should_cancel?
-                        cancel.call
-                        break
-                      end
+                      break if xml_tool_processor.should_cancel?
                     end
-                    partials.each { |inner_partial| blk.call(inner_partial, cancel) }
+                    partials.each { |inner_partial| blk.call(inner_partial) }
                   end
                 end
               end
@@ -248,13 +236,11 @@ def perform_completion!(
                   response_data << stripped
                   result = []
                   result = (xml_tool_processor << stripped) if xml_tool_processor
-                  result.each { |partial| blk.call(partial, cancel) }
+                  result.each { |partial| blk.call(partial) }
                 end
               end
-              if xml_tool_processor
-                xml_tool_processor.finish.each { |partial| blk.call(partial, cancel) }
-              end
-              decode_chunk_finish.each { |partial| blk.call(partial, cancel) }
+              xml_tool_processor.finish.each { |partial| blk.call(partial) } if xml_tool_processor
+              decode_chunk_finish.each { |partial| blk.call(partial) }
               return response_data
             ensure
               if log

lib/personas/bot.rb

@@ -94,7 +94,7 @@ def reply(context, llm_args: {}, &update_blk)
               output_thinking: true,
               cancel_manager: context.cancel_manager,
               **llm_kwargs,
-            ) do |partial, cancel|
+            ) do |partial|
               tool =
                 persona.find_tool(
                   partial,
@@ -111,23 +111,22 @@ def reply(context, llm_args: {}, &update_blk)
                 if tool_call.partial?
                   if tool.class.allow_partial_tool_calls?
                     tool.partial_invoke
-                    update_blk.call("", cancel, tool.custom_raw, :partial_tool)
+                    update_blk.call("", tool.custom_raw, :partial_tool)
                   end
                   next
                 end
 
                 tool_found = true
                 # a bit hacky, but extra newlines do no harm
                 if needs_newlines
-                  update_blk.call("\n\n", cancel)
+                  update_blk.call("\n\n")
                   needs_newlines = false
                 end
 
                 process_tool(
                   tool: tool,
                   raw_context: raw_context,
                   current_llm: current_llm,
-                  cancel: cancel,
                   update_blk: update_blk,
                   prompt: prompt,
                   context: context,
@@ -146,17 +145,17 @@ def reply(context, llm_args: {}, &update_blk)
                 else
                   if partial.is_a?(DiscourseAi::Completions::Thinking)
                     if partial.partial? && partial.message.present?
-                      update_blk.call(partial.message, cancel, nil, :thinking)
+                      update_blk.call(partial.message, nil, :thinking)
                     end
                     if !partial.partial?
                       # this will be dealt with later
                       raw_context << partial
                       current_thinking << partial
                     end
                   elsif partial.is_a?(DiscourseAi::Completions::StructuredOutput)
-                    update_blk.call(partial, cancel, nil, :structured_output)
+                    update_blk.call(partial, nil, :structured_output)
                   else
-                    update_blk.call(partial, cancel)
+                    update_blk.call(partial)
                   end
                 end
               end
@@ -217,14 +216,13 @@ def process_tool(
         tool:,
         raw_context:,
         current_llm:,
-        cancel:,
         update_blk:,
         prompt:,
         context:,
         current_thinking:
       )
         tool_call_id = tool.tool_call_id
-        invocation_result_json = invoke_tool(tool, cancel, context, &update_blk).to_json
+        invocation_result_json = invoke_tool(tool, context, &update_blk).to_json
 
         tool_call_message = {
           type: :tool_call,
@@ -258,27 +256,27 @@ def process_tool(
         raw_context << [invocation_result_json, tool_call_id, "tool", tool.name]
       end
 
-      def invoke_tool(tool, cancel, context, &update_blk)
+      def invoke_tool(tool, context, &update_blk)
         show_placeholder = !context.skip_tool_details && !tool.class.allow_partial_tool_calls?
 
-        update_blk.call("", cancel, build_placeholder(tool.summary, "")) if show_placeholder
+        update_blk.call("", build_placeholder(tool.summary, "")) if show_placeholder
 
         result =
           tool.invoke do |progress, render_raw|
             if render_raw
-              update_blk.call("", cancel, tool.custom_raw, :partial_invoke)
+              update_blk.call("", tool.custom_raw, :partial_invoke)
               show_placeholder = false
             elsif show_placeholder
               placeholder = build_placeholder(tool.summary, progress)
-              update_blk.call("", cancel, placeholder)
+              update_blk.call("", placeholder)
             end
           end
 
         if show_placeholder
           tool_details = build_placeholder(tool.summary, tool.details, custom_raw: tool.custom_raw)
-          update_blk.call(tool_details, cancel, nil, :tool_details)
+          update_blk.call(tool_details, nil, :tool_details)
         elsif tool.custom_raw.present?
-          update_blk.call(tool.custom_raw, cancel, nil, :custom_raw)
+          update_blk.call(tool.custom_raw, nil, :custom_raw)
         end
 
         result

lib/personas/bot_context.rb

@@ -34,7 +34,8 @@ def initialize(
         channel_id: nil,
         context_post_ids: nil,
         feature_name: "bot",
-        resource_url: nil
+        resource_url: nil,
+        cancel_manager: nil
       )
         @participants = participants
         @user = user
@@ -55,6 +56,8 @@ def initialize(
         @feature_name = feature_name
         @resource_url = resource_url
 
+        @cancel_manager = cancel_manager
+
         if post
           @post_id = post.id
           @topic_id = post.topic_id

spec/lib/completions/endpoints/endpoint_compliance.rb

@@ -188,9 +188,11 @@ def streaming_mode_simple_prompt(mock)
     mock.stub_streamed_simple_call(dialect.translate) do
       completion_response = +""
 
-      endpoint.perform_completion!(dialect, user) do |partial, cancel|
+      cancel_manager = DiscourseAi::Completions::CancelManager.new
+
+      endpoint.perform_completion!(dialect, user, cancel_manager: cancel_manager) do |partial|
         completion_response << partial
-        cancel.call if completion_response.split(" ").length == 2
+        cancel_manager.cancel! if completion_response.split(" ").length == 2
       end
 
       expect(AiApiAuditLog.count).to eq(1)
@@ -212,12 +214,14 @@ def streaming_mode_tools(mock)
     prompt = generic_prompt(tools: [mock.tool])
     a_dialect = dialect(prompt: prompt)
 
+    cancel_manager = DiscourseAi::Completions::CancelManager.new
+
     mock.stub_streamed_tool_call(a_dialect.translate) do
       buffered_partial = []
 
-      endpoint.perform_completion!(a_dialect, user) do |partial, cancel|
+      endpoint.perform_completion!(a_dialect, user, cancel_manager: cancel_manager) do |partial|
         buffered_partial << partial
-        cancel.call if partial.is_a?(DiscourseAi::Completions::ToolCall)
+        cancel_manager if partial.is_a?(DiscourseAi::Completions::ToolCall)
       end
 
       expect(buffered_partial).to eq([mock.invocation_response])

spec/lib/modules/ai_bot/playground_spec.rb

@@ -1136,14 +1136,13 @@
 
       split = body.split("|")
 
+      cancel_manager = DiscourseAi::Completions::CancelManager.new
+
       count = 0
       DiscourseAi::AiBot::PostStreamer.on_callback =
         proc do |callback|
           count += 1
-          if count == 2
-            last_post = third_post.topic.posts.order(:id).last
-            Discourse.redis.del("gpt_cancel:#{last_post.id}")
-          end
+          cancel_manager.cancel! if count == 2
           raise "this should not happen" if count > 2
         end
 
@@ -1155,7 +1154,7 @@
         )
         # we are going to need to use real data here cause we want to trigger the
         # base endpoint to cancel part way through
-        playground.reply_to(third_post)
+        playground.reply_to(third_post, cancel_manager: cancel_manager)
       end
 
       last_post = third_post.topic.posts.order(:id).last