From 8b573fe743f4f43391e6610de3daed657d10ea94 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Thu, 3 Apr 2025 00:09:38 +1100
Subject: [PATCH] FIX: maintain newest uploads correctly when constructing
 context (#1242)

---
 lib/completions/prompt_messages_builder.rb    |  5 +-
 .../prompt_messages_builder_spec.rb           | 64 +++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/lib/completions/prompt_messages_builder.rb b/lib/completions/prompt_messages_builder.rb
index b50b5b67..56d2154e 100644
--- a/lib/completions/prompt_messages_builder.rb
+++ b/lib/completions/prompt_messages_builder.rb
@@ -389,8 +389,9 @@ module DiscourseAi
         compressed << current_text if current_text.present?
 
         if upload_count > max_uploads
-          counter = max_uploads - upload_count
-          compressed.delete_if { |item| item.is_a?(Hash) && (counter += 1) > 0 }
+          to_remove = upload_count - max_uploads
+          removed = 0
+          compressed.delete_if { |item| item.is_a?(Hash) && (removed += 1) <= to_remove }
         end
 
         compressed = compressed[0] if compressed.length == 1 && compressed[0].is_a?(String)
diff --git a/spec/lib/completions/prompt_messages_builder_spec.rb b/spec/lib/completions/prompt_messages_builder_spec.rb
index 899b70e4..5fd7eb1d 100644
--- a/spec/lib/completions/prompt_messages_builder_spec.rb
+++ b/spec/lib/completions/prompt_messages_builder_spec.rb
@@ -238,6 +238,70 @@ describe DiscourseAi::Completions::PromptMessagesBuilder do
     end
   end
 
+  describe "upload limits in messages_from_chat" do
+    fab!(:test_channel) { Fabricate(:category_channel) }
+    fab!(:test_user) { Fabricate(:user) }
+
+    # Create MAX_CHAT_UPLOADS + 1 uploads
+    fab!(:uploads) do
+      (described_class::MAX_CHAT_UPLOADS + 1).times.map do |i|
+        Fabricate(:upload, user: test_user, original_filename: "image#{i}.png", extension: "png")
+      end
+    end
+
+    # Create MAX_CHAT_UPLOADS + 1 messages with those uploads
+    fab!(:messages_with_uploads) do
+      uploads.map do |upload|
+        Fabricate(
+          :chat_message,
+          chat_channel: test_channel,
+          user: test_user,
+          message: "Message with upload #{upload.id}",
+        ).tap do |msg|
+          UploadReference.create!(target: msg, upload: upload)
+          msg.update!(upload_ids: [upload.id])
+        end
+      end
+    end
+
+    let(:max_uploads) { described_class::MAX_CHAT_UPLOADS }
+
+    it "limits uploads to MAX_CHAT_UPLOADS in the final result" do
+      last_message = messages_with_uploads.last
+
+      # Make sure uploads are properly associated
+      messages_with_uploads.each_with_index do |msg, i|
+        expect(msg.uploads.first.id).to eq(uploads[i].id)
+      end
+
+      context =
+        described_class.messages_from_chat(
+          last_message,
+          channel: test_channel,
+          context_post_ids: nil,
+          max_messages: messages_with_uploads.size,
+          include_uploads: true,
+          bot_user_ids: [],
+          instruction_message: nil,
+        )
+
+      # We should have one message containing all message content
+      expect(context.length).to eq(1)
+      content = context.first[:content]
+
+      # Count the upload hashes in the content
+      upload_hashes = content.select { |item| item.is_a?(Hash) && item[:upload_id] }
+
+      # Should have exactly MAX_CHAT_UPLOADS upload references
+      expect(upload_hashes.size).to eq(max_uploads)
+
+      # The most recent uploads should be preserved (not the oldest)
+      expected_upload_ids = uploads.last(max_uploads).map(&:id)
+      actual_upload_ids = upload_hashes.map { |h| h[:upload_id] }
+      expect(actual_upload_ids).to match_array(expected_upload_ids)
+    end
+  end
+
   describe ".messages_from_post" do
     fab!(:pm) do
       Fabricate(