DEV: Improve diff streaming with safety checker

Author: keegangeorge

assets/javascripts/discourse/components/modal/diff-modal.gjs

@@ -22,6 +22,7 @@ export default class ModalDiffModal extends Component {
   @service messageBus;
 
   @tracked loading = false;
+  @tracked finalResult = "";
   @tracked diffStreamer = new DiffStreamer(this.args.model.selectedText);
   @tracked suggestion = "";
   @tracked
@@ -65,6 +66,10 @@ export default class ModalDiffModal extends Component {
   async updateResult(result) {
     this.loading = false;
 
+    if (result.done) {
+      this.finalResult = result.result;
+    }
+
     if (this.args.model.showResultAsDiff) {
       this.diffStreamer.updateResult(result, "result");
     } else {
@@ -105,10 +110,14 @@ export default class ModalDiffModal extends Component {
       );
     }
 
-    if (this.args.model.showResultAsDiff && this.diffStreamer.suggestion) {
+    const finalResult =
+      this.finalResult?.length > 0
+        ? this.finalResult
+        : this.diffStreamer.suggestion;
+    if (this.args.model.showResultAsDiff && finalResult) {
       this.args.model.toolbarEvent.replaceText(
         this.args.model.selectedText,
-        this.diffStreamer.suggestion
+        finalResult
       );
     }
   }

lib/ai_helper/assistant.rb

@@ -181,14 +181,16 @@ def stream_prompt(completion_prompt, input, user, channel, force_default_locale:
 
           streamed_diff = parse_diff(input, partial_response) if completion_prompt.diff?
 
-          # Throttle the updates and
-          # checking length prevents partial tags
-          # that aren't sanitized correctly yet (i.e. '<output')
-          #  from being sent in the stream
+          # Throttle updates and check for safe stream points
           if (streamed_result.length > 10 && (Time.now - start > 0.3)) || Rails.env.test?
-            payload = { result: sanitize_result(streamed_result), diff: streamed_diff, done: false }
-            publish_update(channel, payload, user)
-            start = Time.now
+            sanitized = sanitize_result(streamed_result)
+
+            # Skip publishing if inside markdown or HTML link context
+            if DiscourseAi::Utils::DiffUtils::SafetyChecker.safe_to_stream?(sanitized)
+              payload = { result: sanitized, diff: streamed_diff, done: false }
+              publish_update(channel, payload, user)
+              start = Time.now
+            end
           end
         end
 

lib/utils/diff_utils/safety_checker.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require "cgi"
+
+module DiscourseAi
+  module Utils
+    module DiffUtils
+      class SafetyChecker
+        def self.safe_to_stream?(html_text)
+          new(html_text).safe?
+        end
+
+        def initialize(html_text)
+          @original_html = html_text
+          @text = sanitize(html_text)
+        end
+
+        def safe?
+          return false if unclosed_markdown_links?
+          return false if unclosed_raw_html_tag?
+          return false if trailing_incomplete_url?
+          return false if unclosed_backticks?
+          return false if unbalanced_bold_or_italic?
+          return false if incomplete_image_markdown?
+          return false if unbalanced_quote_blocks?
+          return false if unclosed_triple_backticks?
+          return false if partial_emoji?
+
+          true
+        end
+
+        private
+
+        def sanitize(html)
+          text = html.gsub(%r{</?[^>]+>}, "") # remove tags like <span>, <del>, etc.
+          CGI.unescapeHTML(text)
+        end
+
+        def unclosed_markdown_links?
+          open_bracket = @text.rindex("[")
+          close_bracket = @text.rindex("]")
+          open_paren = @text.rindex("(")
+          close_paren = @text.rindex(")")
+          open_bracket && open_paren && (close_bracket.nil? || close_paren.nil?)
+        end
+
+        def unclosed_raw_html_tag?
+          last_lt = @text.rindex("<")
+          last_gt = @text.rindex(">")
+          last_lt && (!last_gt || last_gt < last_lt)
+        end
+
+        def trailing_incomplete_url?
+          last_word = @text.split(/\s/).last
+          last_word =~ %r{\Ahttps?://[^\s]*\z} && last_word !~ /[)\].,!?:;'"]\z/
+        end
+
+        def unclosed_backticks?
+          @text.count("`").odd?
+        end
+
+        def unbalanced_bold_or_italic?
+          @text.scan(/\*\*/).count.odd? || @text.scan(/\*(?!\*)/).count.odd? ||
+            @text.scan(/_/).count.odd?
+        end
+
+        def incomplete_image_markdown?
+          last_image = @text[/!\[.*?\]\(.*?$/, 0]
+          last_image && last_image[-1] != ")"
+        end
+
+        def unbalanced_quote_blocks?
+          opens = @text.scan(/\[quote(=.*?)?\]/i).count
+          closes = @text.scan(%r{\[/quote\]}i).count
+          opens > closes
+        end
+
+        def unclosed_triple_backticks?
+          @text.scan(/```/).count.odd?
+        end
+
+        def partial_emoji?
+          @text
+            .scan(/:[a-z0-9_+.-]*:?/i)
+            .any? do |match|
+              match.count(":") == 1 || (match[-1] != ":" && match =~ /:[a-z0-9_+-]+\.\z/i)
+            end
+        end
+      end
+    end
+  end
+end