FEATURE: allow artifacts to be updated (#980)

Add support for versioned artifacts with improved diff handling * Add versioned artifacts support allowing artifacts to be updated and tracked - New `ai_artifact_versions` table to store version history - Support for updating artifacts through a new `UpdateArtifact` tool - Add version-aware artifact rendering in posts - Include change descriptions for version tracking * Enhance artifact rendering and security - Add support for module-type scripts and external JS dependencies - Expand CSP to allow trusted CDN sources (unpkg, cdnjs, jsdelivr, googleapis) - Improve JavaScript handling in artifacts * Implement robust diff handling system (this is dormant but ready to use once LLMs catch up) - Add new DiffUtils module for applying changes to artifacts - Support for unified diff format with multiple hunks - Intelligent handling of whitespace and line endings - Comprehensive error handling for diff operations * Update routes and UI components - Add versioned artifact routes - Update markdown processing for versioned artifacts Also - Tweaks summary prompt - Improves upload support in custom tool to also provide urls
2025-06-24 08:32:14 +00:00 · 2024-12-03 07:23:31 +11:00 · 2024-12-03 07:23:31 +11:00 · 117c06220e
commit 117c06220e
parent 0ac18d157b
21 changed files with 1007 additions and 32 deletions
--- a/app/controllers/discourse_ai/ai_bot/artifacts_controller.rb
+++ b/app/controllers/discourse_ai/ai_bot/artifacts_controller.rb
@ -19,22 +19,33 @@ module DiscourseAi
          raise Discourse::NotFound if !guardian.can_see?(post)
        end

+        name = artifact.name
+
+        if params[:version].present?
+          artifact = artifact.versions.find_by(version_number: params[:version])
+          raise Discourse::NotFound if !artifact
+        end
+
+        js = artifact.js || ""
+        if !js.match?(%r{\A\s*<script.*</script>}mi)
+          mod = ""
+          mod = " type=\"module\"" if js.match?(/\A\s*import.*/)
+          js = "<script#{mod}>\n#{js}\n</script>"
+        end
        # Prepare the inner (untrusted) HTML document
        untrusted_html = <<~HTML
          <!DOCTYPE html>
          <html>
            <head>
              <meta charset="UTF-8">
-              <title>#{ERB::Util.html_escape(artifact.name)}</title>
+              <title>#{ERB::Util.html_escape(name)}</title>
              <style>
                #{artifact.css}
              </style>
            </head>
            <body>
              #{artifact.html}
-              <script>
-                #{artifact.js}
-              </script>
+              #{js}
            </body>
          </html>
        HTML
@ -45,7 +56,7 @@ module DiscourseAi
          <html>
            <head>
              <meta charset="UTF-8">
-              <title>#{ERB::Util.html_escape(artifact.name)}</title>
+              <title>#{ERB::Util.html_escape(name)}</title>
              <style>
                html, body, iframe {
                  margin: 0;
@ -67,7 +78,9 @@ module DiscourseAi
        HTML

        response.headers.delete("X-Frame-Options")
-        response.headers["Content-Security-Policy"] = "script-src 'unsafe-inline';"
+        response.headers[
+          "Content-Security-Policy"
+        ] = "script-src 'self' 'unsafe-inline' 'wasm-unsafe-eval' https://unpkg.com https://cdnjs.cloudflare.com https://ajax.googleapis.com https://cdn.jsdelivr.net;"
        response.headers["X-Robots-Tag"] = "noindex"

        # Render the content
--- a/app/models/ai_artifact.rb
+++ b/app/models/ai_artifact.rb
@ -1,23 +1,29 @@
 # frozen_string_literal: true

 class AiArtifact < ActiveRecord::Base
+  has_many :versions, class_name: "AiArtifactVersion", dependent: :destroy
  belongs_to :user
  belongs_to :post
  validates :html, length: { maximum: 65_535 }
  validates :css, length: { maximum: 65_535 }
  validates :js, length: { maximum: 65_535 }

-  def self.iframe_for(id)
+  def self.iframe_for(id, version = nil)
    <<~HTML
      <div class='ai-artifact'>
-        <iframe src='#{url(id)}' frameborder="0" height="100%" width="100%"></iframe>
-        <a href='#{url(id)}' target='_blank'>#{I18n.t("discourse_ai.ai_artifact.link")}</a>
+        <iframe src='#{url(id, version)}' frameborder="0" height="100%" width="100%"></iframe>
+        <a href='#{url(id, version)}' target='_blank'>#{I18n.t("discourse_ai.ai_artifact.link")}</a>
      </div>
    HTML
  end

-  def self.url(id)
-    Discourse.base_url + "/discourse-ai/ai-bot/artifacts/#{id}"
+  def self.url(id, version = nil)
+    url = Discourse.base_url + "/discourse-ai/ai-bot/artifacts/#{id}"
+    if version
+      "#{url}/#{version}"
+    else
+      url
+    end
  end

  def self.share_publicly(id:, post:)
@ -33,6 +39,37 @@ class AiArtifact < ActiveRecord::Base
  def url
    self.class.url(id)
  end
+
+  def apply_diff(html_diff: nil, css_diff: nil, js_diff: nil, change_description: nil)
+    differ = DiscourseAi::Utils::DiffUtils
+
+    html = html_diff ? differ.apply_hunk(self.html, html_diff) : self.html
+    css = css_diff ? differ.apply_hunk(self.css, css_diff) : self.css
+    js = js_diff ? differ.apply_hunk(self.js, js_diff) : self.js
+
+    create_new_version(html: html, css: css, js: js, change_description: change_description)
+  end
+
+  def create_new_version(html: nil, css: nil, js: nil, change_description: nil)
+    latest_version = versions.order(version_number: :desc).first
+    new_version_number = latest_version ? latest_version.version_number + 1 : 1
+    version = nil
+
+    transaction do
+      # Create the version record
+      version =
+        versions.create!(
+          version_number: new_version_number,
+          html: html || self.html,
+          css: css || self.css,
+          js: js || self.js,
+          change_description: change_description,
+        )
+      save!
+    end
+
+    version
+  end
 end

 # == Schema Information
--- a/app/models/ai_artifact_version.rb
+++ b/app/models/ai_artifact_version.rb
@ -0,0 +1,27 @@
+# frozen_string_literal: true
+class AiArtifactVersion < ActiveRecord::Base
+  belongs_to :ai_artifact
+  validates :html, length: { maximum: 65_535 }
+  validates :css, length: { maximum: 65_535 }
+  validates :js, length: { maximum: 65_535 }
+end
+
+# == Schema Information
+#
+# Table name: ai_artifact_versions
+#
+#  id                 :bigint           not null, primary key
+#  ai_artifact_id     :bigint           not null
+#  version_number     :integer          not null
+#  html               :string(65535)
+#  css                :string(65535)
+#  js                 :string(65535)
+#  metadata           :jsonb
+#  change_description :string
+#  created_at         :datetime         not null
+#  updated_at         :datetime         not null
+#
+# Indexes
+#
+#  index_ai_artifact_versions_on_ai_artifact_id_and_version_number  (ai_artifact_id,version_number) UNIQUE
+#
--- a/app/models/shared_ai_conversation.rb
+++ b/app/models/shared_ai_conversation.rb
@ -190,9 +190,11 @@ class SharedAiConversation < ActiveRecord::Base
      .css("div.ai-artifact")
      .each do |node|
        id = node["data-ai-artifact-id"].to_i
+        version = node["data-ai-artifact-version"]
+        version_number = version.to_i if version
        if id > 0
          AiArtifact.share_publicly(id: id, post: post)
-          node.replace(AiArtifact.iframe_for(id))
+          node.replace(AiArtifact.iframe_for(id, version_number))
        end
      end

--- a/assets/javascripts/discourse/components/ai-artifact.gjs
+++ b/assets/javascripts/discourse/components/ai-artifact.gjs
@ -37,7 +37,12 @@ export default class AiArtifactComponent extends Component {
  }

  get artifactUrl() {
-    return getURL(`/discourse-ai/ai-bot/artifacts/${this.args.artifactId}`);
+    let url = getURL(`/discourse-ai/ai-bot/artifacts/${this.args.artifactId}`);
+
+    if (this.args.artifactVersion) {
+      url = `${url}/${this.args.artifactVersion}`;
+    }
+    return url;
  }

  @action
--- a/assets/javascripts/initializers/ai-artifacts.gjs
+++ b/assets/javascripts/initializers/ai-artifacts.gjs
@ -14,8 +14,15 @@ function initializeAiArtifacts(api) {
            "data-ai-artifact-id"
          );

+          const artifactVersion = artifactElement.getAttribute(
+            "data-ai-artifact-version"
+          );
+
          helper.renderGlimmer(artifactElement, <template>
-            <AiArtifact @artifactId={{artifactId}} />
+            <AiArtifact
+              @artifactId={{artifactId}}
+              @artifactVersion={{artifactVersion}}
+            />
          </template>);
        }
      );
--- a/assets/javascripts/lib/discourse-markdown/ai-tags.js
+++ b/assets/javascripts/lib/discourse-markdown/ai-tags.js
@ -1,4 +1,8 @@
 export function setup(helper) {
  helper.allowList(["details[class=ai-quote]"]);
-  helper.allowList(["div[class=ai-artifact]", "div[data-ai-artifact-id]"]);
+  helper.allowList([
+    "div[class=ai-artifact]",
+    "div[data-ai-artifact-id]",
+    "div[data-ai-artifact-version]",
+  ]);
 }
--- a/config/locales/server.en.yml
+++ b/config/locales/server.en.yml
@ -205,6 +205,7 @@ en:
    ai_artifact:
      link: "Show Artifact in new tab"
      view_source: "View Source"
+      view_changes: "View Changes"
    unknown_model: "Unknown AI model"

    tools:
@ -309,6 +310,7 @@ en:
            name: "Base Search Query"
            description: "Base query to use when searching. Example: '#urgent' will prepend '#urgent' to the search query and only include topics with the urgent category or tag."
      tool_summary:
+        update_artifact: "Update a web artifact"
        create_artifact: "Create web artifact"
        web_browser: "Browse Web"
        github_search_files: "GitHub search files"
@ -331,6 +333,7 @@ en:
        search_meta_discourse: "Search Meta Discourse"
        javascript_evaluator: "Evaluate JavaScript"
      tool_help:
+        update_artifact: "Update a web artifact using the AI Bot"
        create_artifact: "Create a web artifact using the AI Bot"
        web_browser: "Browse web page using the AI Bot"
        github_search_code: "Search for code in a GitHub repository"
@ -353,6 +356,7 @@ en:
        search_meta_discourse: "Search Meta Discourse"
        javascript_evaluator: "Evaluate JavaScript"
      tool_description:
+        update_artifact: "Updated a web artifact using the AI Bot"
        create_artifact: "Created a web artifact using the AI Bot"
        web_browser: "Reading <a href='%{url}'>%{url}</a>"
        github_search_files: "Searched for '%{keywords}' in %{repo}/%{branch}"
--- a/config/routes.rb
+++ b/config/routes.rb
@ -36,6 +36,7 @@ DiscourseAi::Engine.routes.draw do

  scope module: :ai_bot, path: "/ai-bot/artifacts" do
    get "/:id" => "artifacts#show"
+    get "/:id/:version" => "artifacts#show"
  end

  scope module: :summarization, path: "/summarization", defaults: { format: :json } do
--- a/db/migrate/20241130003808_add_artifact_versions.rb
+++ b/db/migrate/20241130003808_add_artifact_versions.rb
@ -0,0 +1,17 @@
+# frozen_string_literal: true
+class AddArtifactVersions < ActiveRecord::Migration[7.0]
+  def change
+    create_table :ai_artifact_versions do |t|
+      t.bigint :ai_artifact_id, null: false
+      t.integer :version_number, null: false
+      t.string :html, limit: 65_535
+      t.string :css, limit: 65_535
+      t.string :js, limit: 65_535
+      t.jsonb :metadata
+      t.string :change_description
+      t.timestamps
+
+      t.index %i[ai_artifact_id version_number], unique: true
+    end
+  end
+end
--- a/lib/ai_bot/personas/persona.rb
+++ b/lib/ai_bot/personas/persona.rb
@ -99,7 +99,11 @@ module DiscourseAi
              Tools::JavascriptEvaluator,
            ]

-            tools << Tools::CreateArtifact if SiteSetting.ai_artifact_security.in?(%w[lax strict])
+            if SiteSetting.ai_artifact_security.in?(%w[lax strict])
+              tools << Tools::CreateArtifact
+              tools << Tools::UpdateArtifact
+            end
+
            tools << Tools::GithubSearchCode if SiteSetting.ai_bot_github_access_token.present?

            tools << Tools::ListTags if SiteSetting.tagging_enabled
--- a/lib/ai_bot/personas/web_artifact_creator.rb
+++ b/lib/ai_bot/personas/web_artifact_creator.rb
@ -5,11 +5,11 @@ module DiscourseAi
    module Personas
      class WebArtifactCreator < Persona
        def tools
-          [Tools::CreateArtifact]
+          [Tools::CreateArtifact, Tools::UpdateArtifact]
        end

        def required_tools
-          [Tools::CreateArtifact]
+          [Tools::CreateArtifact, Tools::UpdateArtifact]
        end

        def system_prompt
--- a/lib/ai_bot/tool_runner.rb
+++ b/lib/ai_bot/tool_runner.rb
@ -218,7 +218,7 @@ module DiscourseAi
                    for_private_message: @context[:private_message],
                  ).create_for(@bot_user.id)

-                { id: upload.id, short_url: upload.short_url }
+                { id: upload.id, short_url: upload.short_url, url: upload.url }
              end
            ensure
              self.running_attached_function = false
--- a/lib/ai_bot/tools/create_artifact.rb
+++ b/lib/ai_bot/tools/create_artifact.rb
@ -8,6 +8,33 @@ module DiscourseAi
          "create_artifact"
        end

+        def self.js_dependency_tip
+          <<~TIP
+            If you need to include a JavaScript library, you may include assets from:
+            - unpkg.com
+            - cdnjs.com
+            - jsdelivr.com
+            - ajax.googleapis.com
+
+            To include them ensure they are the last tag in your HTML body.
+            Example: <script crossorigin src="https://cdn.jsdelivr.net/npm/vue@2.6.14/dist/vue.min.js"></script>
+          TIP
+        end
+
+        def self.js_script_tag_tip
+          <<~TIP
+            if you need a custom script tag, you can use the following format:
+
+            <script type="module">
+              // your script here
+            </script>
+
+            If you only need a regular script tag, you can use the following format:
+
+            // your script here
+          TIP
+        end
+
        def self.signature
          {
            name: "create_artifact",
@ -22,7 +49,8 @@ module DiscourseAi
              },
              {
                name: "html_body",
-                description: "The HTML content for the BODY tag (do not include the BODY tag)",
+                description:
+                  "The HTML content for the BODY tag (do not include the BODY tag). #{js_dependency_tip}",
                type: "string",
                required: true,
              },
@ -31,7 +59,7 @@ module DiscourseAi
                name: "js",
                description:
                  "Optional
-              JavaScript code for the artifact",
+              JavaScript code for the artifact, this will be the last <script> tag in the BODY. #{js_script_tag_tip}",
                type: "string",
              },
            ],
--- a/lib/ai_bot/tools/update_artifact.rb
+++ b/lib/ai_bot/tools/update_artifact.rb
@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module AiBot
+    module Tools
+      class UpdateArtifact < Tool
+        def self.name
+          "update_artifact"
+        end
+
+        # this is not working that well, we support it, but I am leaving it dormant for now
+        def self.unified_diff_tip
+          <<~TIP
+            When updating and artifact in diff mode unified diffs can be applied:
+
+            If editing:
+
+            <div>
+              <p>Some text</p>
+            </div>
+
+            You can provide a diff like:
+
+            <div>
+            - <p>Some text</p>
+            + <p>Some new text</p>
+            </div>
+
+            This will result in:
+
+            <div>
+              <p>Some new text</p>
+            </div>
+
+            If you need to supply multiple hunks for a diff use a @@ separator, for example:
+
+            @@ -1,3 +1,3 @@
+            - <p>Some text</p>
+            + <p>Some new text</p>
+            @@ -5,3 +5,3 @@
+            - </div>
+            + <p>more text</p>
+            </div>
+
+            If you supply text without @@ seperators or + and - prefixes, the entire text will be appended to the artifact section.
+
+          TIP
+        end
+
+        def self.signature
+          {
+            name: "update_artifact",
+            description:
+              "Updates an existing web artifact with new HTML, CSS, or JavaScript content. Note either html, css, or js MUST be provided. You may provide all three if desired.",
+            parameters: [
+              {
+                name: "artifact_id",
+                description: "The ID of the artifact to update",
+                type: "integer",
+                required: true,
+              },
+              { name: "html", description: "new HTML content for the artifact", type: "string" },
+              { name: "css", description: "new CSS content for the artifact", type: "string" },
+              {
+                name: "js",
+                description: "new JavaScript content for the artifact",
+                type: "string",
+              },
+              {
+                name: "change_description",
+                description:
+                  "A brief description of the changes being made. Note: This only documents the change - you must provide the actual content in html/css/js parameters to make changes.",
+                type: "string",
+              },
+            ],
+          }
+        end
+
+        def self.allow_partial_tool_calls?
+          true
+        end
+
+        def chain_next_response?
+          @chain_next_response
+        end
+
+        def partial_invoke
+          @selected_tab = :html
+          if @prev_parameters
+            @selected_tab = parameters.keys.find { |k| @prev_parameters[k] != parameters[k] }
+          end
+          update_custom_html
+          @prev_parameters = parameters.dup
+        end
+
+        def invoke
+          yield "Updating Artifact"
+
+          post = Post.find_by(id: context[:post_id])
+          return error_response("No post context found") unless post
+
+          artifact = AiArtifact.find_by(id: parameters[:artifact_id])
+          return error_response("Artifact not found") unless artifact
+
+          if artifact.post.topic.id != post.topic.id
+            return error_response("Attempting to update an artifact you are not allowed to")
+          end
+
+          last_version = artifact.versions.order(version_number: :desc).first
+
+          begin
+            version =
+              artifact.create_new_version(
+                html: parameters[:html] || last_version&.html || artifact.html,
+                css: parameters[:css] || last_version&.css || artifact.css,
+                js: parameters[:js] || last_version&.js || artifact.js,
+                change_description: parameters[:change_description].to_s,
+              )
+
+            update_custom_html(artifact, version)
+            success_response(artifact, version)
+          rescue DiscourseAi::Utils::DiffUtils::DiffError => e
+            error_response(e.to_llm_message)
+          rescue => e
+            error_response(e.message)
+          end
+        end
+
+        private
+
+        def update_custom_html(artifact = nil, version = nil)
+          content = []
+
+          if parameters[:html].present?
+            content << [:html, "### HTML Changes\n\n```html\n#{parameters[:html]}\n```"]
+          end
+
+          if parameters[:css].present?
+            content << [:css, "### CSS Changes\n\n```css\n#{parameters[:css]}\n```"]
+          end
+
+          if parameters[:js].present?
+            content << [:js, "### JavaScript Changes\n\n```javascript\n#{parameters[:js]}\n```"]
+          end
+
+          if parameters[:change_description].present?
+            content.unshift(
+              [:description, "### Change Description\n\n#{parameters[:change_description]}"],
+            )
+          end
+
+          content.sort_by! { |c| c[0] === @selected_tab ? 1 : 0 } if !artifact
+
+          if artifact
+            content.unshift([nil, "[details='#{I18n.t("discourse_ai.ai_artifact.view_changes")}']"])
+            content << [nil, "[/details]"]
+            content << [
+              :preview,
+              "### Preview\n\n<div class=\"ai-artifact\" data-ai-artifact-version=\"#{version.version_number}\" data-ai-artifact-id=\"#{artifact.id}\"></div>",
+            ]
+          end
+
+          content.unshift("\n\n")
+
+          self.custom_raw = content.map { |c| c[1] }.join("\n\n")
+        end
+
+        def success_response(artifact, version)
+          @chain_next_response = false
+
+          hash = {
+            status: "success",
+            artifact_id: artifact.id,
+            version: version.version_number,
+            message: "Artifact updated successfully and rendered to user.",
+          }
+
+          hash
+        end
+
+        def error_response(message)
+          @chain_next_response = true
+          self.custom_raw = ""
+
+          { status: "error", error: message }
+        end
+
+        def help
+          "Updates an existing web artifact with changes to its HTML, CSS, or JavaScript content. " \
+            "Requires the artifact ID and at least one change diff. " \
+            "Changes are applied using unified diff format. " \
+            "A description of the changes is required for version history."
+        end
+      end
+    end
+  end
+end
--- a/lib/summarization/strategies/topic_summary.rb
+++ b/lib/summarization/strategies/topic_summary.rb
@ -42,10 +42,10 @@ module DiscourseAi
            - Maintain the original language of the text being summarized.
            - Aim for summaries to be 400 words or less.
            - Each new post is formatted as "<POST_NUMBER>) <USERNAME> <MESSAGE>"
-            - Cite specific noteworthy posts using the format [NAME](#{resource_path}/POST_NUMBER)
-              - Example: link to the 3rd post by sam: [sam](#{resource_path}/3)
+            - Cite specific noteworthy posts using the format [DESCRIPTION](#{resource_path}/POST_NUMBER)
+              - Example: links to the 3rd and 6th posts by sam: sam ([#3](#{resource_path}/3), [#6](#{resource_path}/6))
              - Example: link to the 6th post by jane: [agreed with](#{resource_path}/6)
-              - Example: link to the 13th post by joe: [#13](#{resource_path}/13)
+              - Example: link to the 13th post by joe: [joe](#{resource_path}/13)
            - When formatting usernames either use @USERNAME or [USERNAME](#{resource_path}/POST_NUMBER)
          TEXT

@ -53,11 +53,11 @@ module DiscourseAi
            ### Context:

            #{content_title.present? ? "The discussion title is: " + content_title + ".\n" : ""}
-              
+
            Here is the existing summary:
-              
+
            #{summary}
-          
+
            Here are the new posts, inside <input></input> XML tags:

            <input>
@ -84,10 +84,10 @@ module DiscourseAi
            - Maintain the original language of the text being summarized.
            - Aim for summaries to be 400 words or less.
            - Each post is formatted as "<POST_NUMBER>) <USERNAME> <MESSAGE>"
-            - Cite specific noteworthy posts using the format [NAME](#{resource_path}/POST_NUMBER)
-              - Example: link to the 3rd post by sam: [sam](#{resource_path}/3)
+            - Cite specific noteworthy posts using the format [DESCRIPTION](#{resource_path}/POST_NUMBER)
+              - Example: links to the 3rd and 6th posts by sam: sam ([#3](#{resource_path}/3), [#6](#{resource_path}/6))
              - Example: link to the 6th post by jane: [agreed with](#{resource_path}/6)
-              - Example: link to the 13th post by joe: [#13](#{resource_path}/13)
+              - Example: link to the 13th post by joe: [joe](#{resource_path}/13)
            - When formatting usernames either use @USERNMAE OR [USERNAME](#{resource_path}/POST_NUMBER)
          TEXT

--- a/lib/utils/diff_utils.rb
+++ b/lib/utils/diff_utils.rb
@ -0,0 +1,184 @@
+# frozen_string_literal: true
+# Inspired by Aider https://github.com/Aider-AI/aider
+
+module DiscourseAi
+  module Utils
+    module DiffUtils
+      # Custom errors with detailed information for LLM feedback
+      class DiffError < StandardError
+        attr_reader :original_text, :diff_text, :context
+
+        def initialize(message, original_text:, diff_text:, context: {})
+          @original_text = original_text
+          @diff_text = diff_text
+          @context = context
+          super(message)
+        end
+
+        def to_llm_message
+          original_text = @original_text
+          original_text = @original_text[0..1000] + "..." if @original_text.length > 1000
+
+          <<~MESSAGE
+            #{message}
+
+            Original text:
+            ```
+            #{original_text}
+            ```
+
+            Attempted diff:
+            ```
+            #{diff_text}
+            ```
+
+            #{context_message}
+
+            Please provide a corrected diff that:
+            1. Has the correct context lines
+            2. Contains all necessary removals (-) and additions (+)
+          MESSAGE
+        end
+
+        private
+
+        def context_message
+          return "" if context.empty?
+
+          context.map { |key, value| "#{key}: #{value}" }.join("\n")
+        end
+      end
+
+      class NoMatchingContextError < DiffError
+        def initialize(original_text:, diff_text:)
+          super(
+            "Could not find the context lines in the original text",
+            original_text: original_text,
+            diff_text: diff_text,
+          )
+        end
+      end
+
+      class AmbiguousMatchError < DiffError
+        def initialize(original_text:, diff_text:)
+          super(
+            "Found multiple possible locations for this change",
+            original_text: original_text,
+            diff_text: diff_text,
+          )
+        end
+      end
+
+      class MalformedDiffError < DiffError
+        def initialize(original_text:, diff_text:, issue:)
+          super(
+            "The diff format is invalid",
+            original_text: original_text,
+            diff_text: diff_text,
+            context: {
+              "Issue" => issue,
+            },
+          )
+        end
+      end
+
+      def self.apply_hunk(text, diff)
+        # we need to handle multiple hunks just in case
+        if diff.match?(/^\@\@.*\@\@$\n/)
+          hunks = diff.split(/^\@\@.*\@\@$\n/)
+          if hunks.present?
+            hunks.each do |hunk|
+              next if hunk.blank?
+              text = apply_hunk(text, hunk)
+            end
+            return text
+          end
+        end
+
+        text = text.encode(universal_newline: true)
+        diff = diff.encode(universal_newline: true)
+        # we need this for matching
+        text = text + "\n" unless text.end_with?("\n")
+
+        diff_lines = parse_diff_lines(diff, text)
+
+        validate_diff_format!(text, diff, diff_lines)
+
+        return text.strip + "\n" + diff.strip if diff_lines.all? { |marker, _| marker == " " }
+
+        lines_to_match = diff_lines.select { |marker, _| ["-", " "].include?(marker) }.map(&:last)
+        match_start, match_end = find_unique_match(text, lines_to_match, diff)
+        new_hunk = diff_lines.select { |marker, _| ["+", " "].include?(marker) }.map(&:last).join
+
+        new_hunk = +""
+
+        diff_lines_index = 0
+        text[match_start..match_end].lines.each do |line|
+          diff_marker, diff_content = diff_lines[diff_lines_index]
+
+          while diff_marker == "+"
+            new_hunk << diff_content
+            diff_lines_index += 1
+            diff_marker, diff_content = diff_lines[diff_lines_index]
+          end
+
+          new_hunk << line if diff_marker == " "
+
+          diff_lines_index += 1
+        end
+
+        # leftover additions
+        diff_marker, diff_content = diff_lines[diff_lines_index]
+        while diff_marker == "+"
+          diff_lines_index += 1
+          new_hunk << diff_content
+          diff_marker, diff_content = diff_lines[diff_lines_index]
+        end
+
+        (text[0...match_start].to_s + new_hunk + text[match_end..-1].to_s).strip
+      end
+
+      private_class_method def self.parse_diff_lines(diff, text)
+        diff.lines.map do |line|
+          marker = line[0]
+          content = line[1..]
+
+          if !["-", "+", " "].include?(marker)
+            marker = " "
+            content = line
+          end
+
+          [marker, content]
+        end
+      end
+
+      private_class_method def self.validate_diff_format!(text, diff, diff_lines)
+        if diff_lines.empty?
+          raise MalformedDiffError.new(original_text: text, diff_text: diff, issue: "Diff is empty")
+        end
+      end
+
+      private_class_method def self.find_unique_match(text, context_lines, diff)
+        return 0 if context_lines.empty? && removals.empty?
+
+        pattern = context_lines.map { |line| "^\\s*" + Regexp.escape(line.strip) + "\s*$\n" }.join
+        matches =
+          text
+            .enum_for(:scan, /#{pattern}/m)
+            .map do
+              match = Regexp.last_match
+              [match.begin(0), match.end(0)]
+            end
+
+        case matches.length
+        when 0
+          raise NoMatchingContextError.new(original_text: text, diff_text: diff)
+        when 1
+          matches.first
+        else
+          raise AmbiguousMatchError.new(original_text: text, diff_text: diff)
+        end
+      end
+    end
+  end
+end
--- a/lib/utils/dns_srv.rb
+++ b/lib/utils/dns_srv.rb
@ -19,7 +19,7 @@ module DiscourseAi

      def self.dns_srv_lookup_for_domain(domain)
        resolver = Resolv::DNS.new
-        resources = resolver.getresources(domain, Resolv::DNS::Resource::IN::SRV)
+        resolver.getresources(domain, Resolv::DNS::Resource::IN::SRV)
      end

      def self.select_server(resources)
--- a/spec/lib/modules/ai_bot/tools/update_artifact_spec.rb
+++ b/spec/lib/modules/ai_bot/tools/update_artifact_spec.rb
@ -0,0 +1,148 @@
+# spec/lib/modules/ai_bot/tools/update_artifact_spec.rb
+# frozen_string_literal: true
+
+RSpec.describe DiscourseAi::AiBot::Tools::UpdateArtifact do
+  fab!(:llm_model)
+  let(:bot_user) { DiscourseAi::AiBot::EntryPoint.find_user_from_model(llm_model.name) }
+  let(:llm) { DiscourseAi::Completions::Llm.proxy("custom:#{llm_model.id}") }
+  fab!(:post)
+  fab!(:artifact) do
+    AiArtifact.create!(
+      user: Fabricate(:user),
+      post: post,
+      name: "Test Artifact",
+      html: "<div>\nOriginal\n</div>",
+      css: "div {\n color: blue; \n}",
+      js: "console.log('hello');",
+    )
+  end
+
+  before { SiteSetting.ai_bot_enabled = true }
+
+  describe "#process" do
+    let(:html) { <<~DIFF }
+         <div>
+           Updated
+         </div>
+      DIFF
+
+    let(:css) { <<~DIFF }
+         div {
+          color: red;
+         }
+      DIFF
+
+    let(:js) { <<~DIFF }
+        console.log('world');
+      DIFF
+
+    it "updates artifact content when supplied" do
+      tool =
+        described_class.new(
+          {
+            artifact_id: artifact.id,
+            html: html,
+            css: css,
+            js: js,
+            change_description: "Updated colors and text",
+          },
+          bot_user: bot_user,
+          llm: llm,
+          context: {
+            post_id: post.id,
+          },
+        )
+
+      result = tool.invoke {}
+
+      expect(result[:status]).to eq("success")
+      expect(result[:version]).to eq(1)
+
+      version = artifact.versions.find_by(version_number: 1)
+      expect(version.html).to include("Updated")
+      expect(version.css).to include("color: red")
+      expect(version.js).to include("'world'")
+      expect(artifact.versions.count).to eq(1)
+      expect(version.change_description).to eq("Updated colors and text")
+
+      # updating again should update the correct version
+      tool.parameters = {
+        artifact_id: artifact.id,
+        html: nil,
+        css: nil,
+        js: "updated",
+        change_description: "Updated colors and text again",
+      }
+
+      result = tool.invoke {}
+
+      version = artifact.versions.find_by(version_number: 2)
+
+      expect(result[:status]).to eq("success")
+      expect(result[:version]).to eq(2)
+
+      expect(version.html).to include("Updated")
+      expect(version.css).to include("color: red")
+      expect(version.js).to include("updated")
+    end
+
+    it "handles partial updates correctly" do
+      tool = described_class.new({}, bot_user: bot_user, llm: llm)
+
+      tool.parameters = { artifact_id: artifact.id, html: html, change_description: "Changed HTML" }
+      tool.partial_invoke
+
+      expect(tool.custom_raw).to include("### HTML Changes")
+      expect(tool.custom_raw).to include("### Change Description")
+      expect(tool.custom_raw).not_to include("### CSS Changes")
+    end
+
+    it "handles invalid artifact ID" do
+      tool =
+        described_class.new(
+          { artifact_id: -1, html: html, change_description: "Test change" },
+          bot_user: bot_user,
+          llm: llm,
+          context: {
+            post_id: post.id,
+          },
+        )
+
+      result = tool.invoke {}
+      expect(result[:status]).to eq("error")
+      expect(result[:error]).to eq("Artifact not found")
+    end
+
+    it "requires at least one change" do
+      tool =
+        described_class.new(
+          { artifact_id: artifact.id, change_description: "No changes" },
+          bot_user: bot_user,
+          llm: llm,
+          context: {
+            post_id: post.id,
+          },
+        )
+
+      result = tool.invoke {}
+      expect(result[:status]).to eq("success")
+      expect(artifact.versions.count).to eq(1)
+    end
+
+    it "correctly renders changes in message" do
+      tool =
+        described_class.new(
+          { artifact_id: artifact.id, html: html, change_description: "Updated content" },
+          bot_user: bot_user,
+          llm: llm,
+          context: {
+            post_id: post.id,
+          },
+        )
+
+      tool.invoke {}
+
+      expect(tool.custom_raw.strip).to include(html.strip)
+    end
+  end
+end
--- a/spec/lib/utils/diff_utils_spec.rb
+++ b/spec/lib/utils/diff_utils_spec.rb
@ -0,0 +1,284 @@
+# frozen_string_literal: true
+
+RSpec.describe DiscourseAi::Utils::DiffUtils do
+  describe ".apply_hunk" do
+    subject(:apply_hunk) { described_class.apply_hunk(original_text, diff) }
+
+    context "with HTML content" do
+      let(:original_text) { <<~HTML }
+          <div class="container">
+            <h1>Original Title</h1>
+            <p>Some content here</p>
+            <ul>
+              <li>Item 1</li>
+              <li>Item 2</li>
+            </ul>
+          </div>
+        HTML
+
+      context "when adding content" do
+        let(:diff) { <<~DIFF }
+            <div class="container">
+              <h1>Original Title</h1>
+            +  <h2>New Subtitle</h2>
+              <p>Some content here</p>
+          DIFF
+
+        it "inserts the new content" do
+          expected = <<~HTML
+            <div class="container">
+              <h1>Original Title</h1>
+              <h2>New Subtitle</h2>
+              <p>Some content here</p>
+              <ul>
+                <li>Item 1</li>
+                <li>Item 2</li>
+              </ul>
+            </div>
+          HTML
+          expect(apply_hunk).to eq(expected.strip)
+        end
+      end
+
+      context "when removing content" do
+        let(:diff) { <<~DIFF }
+            <ul>
+            - <li>Item 1</li>
+            <li>Item 2</li>
+          DIFF
+
+        it "removes the specified content" do
+          # note how this is super forgiving
+          expected = <<~HTML
+            <div class="container">
+              <h1>Original Title</h1>
+              <p>Some content here</p>
+              <ul>
+                <li>Item 2</li>
+              </ul>
+            </div>
+          HTML
+          expect(apply_hunk).to eq(expected.strip)
+        end
+      end
+
+      context "when replacing content" do
+        let(:diff) { <<~DIFF }
+             <div class="container">
+            -  <h1>Original Title</h1>
+            +  <h1>Updated Title</h1>
+               <p>Some content here</p>
+          DIFF
+
+        it "replaces the content correctly" do
+          expected = <<~HTML
+            <div class="container">
+              <h1>Updated Title</h1>
+              <p>Some content here</p>
+              <ul>
+                <li>Item 1</li>
+                <li>Item 2</li>
+              </ul>
+            </div>
+          HTML
+
+          expect(apply_hunk).to eq(expected.strip)
+        end
+      end
+    end
+
+    context "with CSS content" do
+      let(:original_text) { <<~CSS }
+          .container {
+            background: #fff;
+            padding: 20px;
+            margin: 10px;
+          }
+        CSS
+
+      context "when modifying properties" do
+        let(:diff) { <<~DIFF }
+             .container {
+            -  background: #fff;
+            +  background: #f5f5f5;
+            +  happy: sam;
+            -  padding: 20px;
+            +  padding: 10px;
+          DIFF
+
+        it "updates the property value" do
+          expected = <<~CSS
+            .container {
+              background: #f5f5f5;
+              happy: sam;
+              padding: 10px;
+              margin: 10px;
+            }
+          CSS
+
+          expect(apply_hunk).to eq(expected.strip)
+        end
+      end
+    end
+
+    context "when handling errors" do
+      let(:original_text) { <<~HTML }
+          <div>
+            <h1>Title</h1>
+            <p>
+            <h1>Title</h1>
+            <p>
+          </div>
+        HTML
+
+      context "with ambiguous matches" do
+        let(:diff) { <<~DIFF }
+             <h1>Title</h1>
+            +<h2>Subtitle</h2>
+             <p>
+          DIFF
+
+        it "raises an AmbiguousMatchError" do
+          expect { apply_hunk }.to raise_error(
+            DiscourseAi::Utils::DiffUtils::AmbiguousMatchError,
+          ) do |error|
+            expect(error.to_llm_message).to include("Found multiple possible locations")
+          end
+        end
+      end
+
+      context "with no matching context" do
+        let(:diff) { <<~DIFF }
+             <h1>Wrong Title</h1>
+            +<h2>Subtitle</h2>
+             <p>
+          DIFF
+
+        it "raises a NoMatchingContextError" do
+          expect { apply_hunk }.to raise_error(
+            DiscourseAi::Utils::DiffUtils::NoMatchingContextError,
+          ) do |error|
+            expect(error.to_llm_message).to include("Could not find the context lines")
+          end
+        end
+      end
+
+      context "with malformed diffs" do
+        context "when empty" do
+          let(:diff) { "" }
+
+          it "raises a MalformedDiffError" do
+            expect { apply_hunk }.to raise_error(
+              DiscourseAi::Utils::DiffUtils::MalformedDiffError,
+            ) do |error|
+              expect(error.context["Issue"]).to eq("Diff is empty")
+            end
+          end
+        end
+      end
+    end
+
+    context "without markers" do
+      let(:original_text) { "hello" }
+      let(:diff) { "world" }
+      it "will append to the end" do
+        expect(apply_hunk).to eq("hello\nworld")
+      end
+    end
+
+    context "when appending text to the end of a document" do
+      let(:original_text) { "hello\nworld" }
+
+      let(:diff) { <<~DIFF }
+          world
+          +123
+        DIFF
+
+      it "can append to end" do
+        expect(apply_hunk).to eq("hello\nworld\n123")
+      end
+    end
+
+    context "when applying multiple hunks to a file" do
+      let(:original_text) { <<~TEXT }
+          1
+          2
+          3
+          4
+          5
+          6
+          7
+          8
+        TEXT
+
+      let(:diff) { <<~DIFF }
+          @@ -1,4 +1,4 @@
+          2
+          - 3
+          @@ -6,4 +6,4 @@
+          - 7
+        DIFF
+
+      it "can apply multiple hunks" do
+        expected = <<~TEXT
+          1
+          2
+          4
+          5
+          6
+          8
+        TEXT
+        expect(apply_hunk).to eq(expected.strip)
+      end
+    end
+
+    context "with line ending variations" do
+      let(:original_text) { "line1\r\nline2\nline3\r\n" }
+      let(:diff) { <<~DIFF }
+           line1
+          +new line
+           line2
+        DIFF
+
+      it "handles mixed line endings" do
+        expect(apply_hunk).to include("new line")
+        expect(apply_hunk.lines.count).to eq(4)
+      end
+    end
+
+    context "with whitespace sensitivity" do
+      let(:original_text) { <<~TEXT }
+          def method
+              puts "hello"
+          end
+        TEXT
+
+      context "when indentation matters" do
+        let(:diff) { <<~DIFF }
+             def method
+            -    puts "hello"
+            +    puts "world"
+               end
+          DIFF
+
+        it "preserves exact indentation" do
+          result = apply_hunk
+          expect(result).to match(/^    puts "world"$/)
+        end
+      end
+
+      context "when trailing whitespace exists" do
+        let(:original_text) { "line1  \nline2\n" }
+        let(:diff) { <<~DIFF }
+             line1
+            +new line
+             line2
+          DIFF
+
+        it "preserves significant whitespace" do
+          expect(apply_hunk).to include("line1  \n")
+        end
+      end
+    end
+  end
+end
--- a/spec/requests/ai_bot/artifacts_controller_spec.rb
+++ b/spec/requests/ai_bot/artifacts_controller_spec.rb
@ -55,6 +55,19 @@ RSpec.describe DiscourseAi::AiBot::ArtifactsController do
        expect(untrusted_html).to include(artifact.css)
        expect(untrusted_html).to include(artifact.js)
      end
+
+      it "can also find an artifact by its version" do
+        sign_in(user)
+
+        version = artifact.create_new_version(html: "<div>Was Updated</div>")
+
+        get "/discourse-ai/ai-bot/artifacts/#{artifact.id}/#{version.version_number}"
+        expect(response.status).to eq(200)
+        untrusted_html = parse_srcdoc(response.body)
+        expect(untrusted_html).to include("Was Updated")
+        expect(untrusted_html).to include(artifact.css)
+        expect(untrusted_html).to include(artifact.js)
+      end
    end

    context "with public artifact" do
@ -71,7 +84,7 @@ RSpec.describe DiscourseAi::AiBot::ArtifactsController do
      sign_in(user)
      get "/discourse-ai/ai-bot/artifacts/#{artifact.id}"
      expect(response.headers["X-Frame-Options"]).to eq(nil)
-      expect(response.headers["Content-Security-Policy"]).to eq("script-src 'unsafe-inline';")
+      expect(response.headers["Content-Security-Policy"]).to include("unsafe-inline")
      expect(response.headers["X-Robots-Tag"]).to eq("noindex")
    end
  end