diff --git a/app/controllers/discourse_ai/admin/ai_tools_controller.rb b/app/controllers/discourse_ai/admin/ai_tools_controller.rb
index f5aa081b..92f5905c 100644
--- a/app/controllers/discourse_ai/admin/ai_tools_controller.rb
+++ b/app/controllers/discourse_ai/admin/ai_tools_controller.rb
@@ -5,7 +5,7 @@ module DiscourseAi
     class AiToolsController < ::Admin::AdminController
       requires_plugin ::DiscourseAi::PLUGIN_NAME
 
-      before_action :find_ai_tool, only: %i[show update destroy]
+      before_action :find_ai_tool, only: %i[test show update destroy]
 
       def index
         ai_tools = AiTool.all
@@ -17,7 +17,7 @@ module DiscourseAi
       end
 
       def create
-        ai_tool = AiTool.new(ai_tool_params.except(:rag_uploads))
+        ai_tool = AiTool.new(ai_tool_params)
         ai_tool.created_by_id = current_user.id
 
         if ai_tool.save
@@ -29,7 +29,7 @@ module DiscourseAi
       end
 
       def update
-        if @ai_tool.update(ai_tool_params.except(:rag_uploads))
+        if @ai_tool.update(ai_tool_params)
           RagDocumentFragment.update_target_uploads(@ai_tool, attached_upload_ids)
           render_serialized(@ai_tool, AiCustomToolSerializer)
         else
@@ -46,18 +46,13 @@ module DiscourseAi
       end
 
       def test
-        if params[:id].present?
-          ai_tool = AiTool.find(params[:id])
-        else
-          ai_tool = AiTool.new(ai_tool_params.except(:rag_uploads))
-        end
-
+        @ai_tool.assign_attributes(ai_tool_params) if params[:ai_tool]
         parameters = params[:parameters].to_unsafe_h
 
         # we need an llm so we have a tokenizer
         # but will do without if none is available
         llm = LlmModel.first&.to_llm
-        runner = ai_tool.runner(parameters, llm: llm, bot_user: current_user, context: {})
+        runner = @ai_tool.runner(parameters, llm: llm, bot_user: current_user, context: {})
         result = runner.invoke
 
         if result.is_a?(Hash) && result[:error]
@@ -74,24 +69,27 @@ module DiscourseAi
       private
 
       def attached_upload_ids
-        ai_tool_params[:rag_uploads].to_a.map { |h| h[:id] }
+        params[:ai_tool][:rag_uploads].to_a.map { |h| h[:id] }
       end
 
       def find_ai_tool
-        @ai_tool = AiTool.find(params[:id])
+        @ai_tool = AiTool.find(params[:id].to_i)
       end
 
       def ai_tool_params
-        params.require(:ai_tool).permit(
-          :name,
-          :description,
-          :script,
-          :summary,
-          :rag_chunk_tokens,
-          :rag_chunk_overlap_tokens,
-          rag_uploads: [:id],
-          parameters: [:name, :type, :description, :required, enum: []],
-        )
+        params
+          .require(:ai_tool)
+          .permit(
+            :name,
+            :description,
+            :script,
+            :summary,
+            :rag_chunk_tokens,
+            :rag_chunk_overlap_tokens,
+            rag_uploads: [:id],
+            parameters: [:name, :type, :description, :required, enum: []],
+          )
+          .except(:rag_uploads)
       end
     end
   end
diff --git a/assets/javascripts/discourse/components/ai-tool-editor.gjs b/assets/javascripts/discourse/components/ai-tool-editor.gjs
index 566b704a..24536e62 100644
--- a/assets/javascripts/discourse/components/ai-tool-editor.gjs
+++ b/assets/javascripts/discourse/components/ai-tool-editor.gjs
@@ -232,11 +232,13 @@ export default class AiToolEditor extends Component {
         {{/if}}
 
         <div class="control-group ai-tool-editor__action_panel">
-          <DButton
-            @action={{this.openTestModal}}
-            @label="discourse_ai.tools.test"
-            class="ai-tool-editor__test-button"
-          />
+          {{#unless @model.isNew}}
+            <DButton
+              @action={{this.openTestModal}}
+              @label="discourse_ai.tools.test"
+              class="ai-tool-editor__test-button"
+            />
+          {{/unless}}
 
           <DButton
             @action={{this.save}}
diff --git a/assets/javascripts/discourse/components/modal/ai-tool-test-modal.gjs b/assets/javascripts/discourse/components/modal/ai-tool-test-modal.gjs
index a188a2e6..7df6b32a 100644
--- a/assets/javascripts/discourse/components/modal/ai-tool-test-modal.gjs
+++ b/assets/javascripts/discourse/components/modal/ai-tool-test-modal.gjs
@@ -25,7 +25,7 @@ export default class AiToolTestModal extends Component {
     this.isLoading = true;
     try {
       const response = await ajax(
-        "/admin/plugins/discourse-ai/ai-tools/test.json",
+        `/admin/plugins/discourse-ai/ai-tools/${this.args.model.tool.id}/test.json`,
         {
           type: "POST",
           data: JSON.stringify({
diff --git a/config/routes.rb b/config/routes.rb
index 461e8b4b..e806f741 100644
--- a/config/routes.rb
+++ b/config/routes.rb
@@ -55,7 +55,9 @@ Discourse::Application.routes.draw do
       only: %i[index create show update destroy],
       path: "ai-tools",
       controller: "discourse_ai/admin/ai_tools",
-    ) { post :test, on: :collection }
+    )
+
+    post "/ai-tools/:id/test", to: "discourse_ai/admin/ai_tools#test"
 
     post "/ai-personas/:id/create-user", to: "discourse_ai/admin/ai_personas#create_user"
 
diff --git a/lib/ai_bot/tool_runner.rb b/lib/ai_bot/tool_runner.rb
index 21c5d68f..67adeed1 100644
--- a/lib/ai_bot/tool_runner.rb
+++ b/lib/ai_bot/tool_runner.rb
@@ -149,6 +149,7 @@ module DiscourseAi
             limit: limit,
             offset: 0,
           )
+
         fragments =
           RagDocumentFragment.where(id: fragment_ids, upload_id: upload_refs).pluck(
             :id,
@@ -174,8 +175,9 @@ module DiscourseAi
       def attach_index(mini_racer_context)
         mini_racer_context.attach(
           "_index_search",
-          ->(query, options) do
+          ->(*params) do
             begin
+              query, options = params
               self.running_attached_function = true
               options ||= {}
               options = options.symbolize_keys
diff --git a/lib/embeddings/vector_representations/base.rb b/lib/embeddings/vector_representations/base.rb
index d8236181..d5c23d21 100644
--- a/lib/embeddings/vector_representations/base.rb
+++ b/lib/embeddings/vector_representations/base.rb
@@ -205,7 +205,10 @@ module DiscourseAi
                   FROM
                     #{rag_fragments_table_name}
                   INNER JOIN
-                    rag_document_fragments ON rag_document_fragments.id = rag_document_fragment_id
+                    rag_document_fragments ON
+                      rag_document_fragments.id = rag_document_fragment_id AND
+                      rag_document_fragments.target_id = :target_id AND
+                      rag_document_fragments.target_type = :target_type
                   WHERE
                     model_id = #{id} AND strategy_id = #{@strategy.id}
                   ORDER BY
diff --git a/spec/requests/admin/ai_tools_controller_spec.rb b/spec/requests/admin/ai_tools_controller_spec.rb
index e05b7344..3ba96ac4 100644
--- a/spec/requests/admin/ai_tools_controller_spec.rb
+++ b/spec/requests/admin/ai_tools_controller_spec.rb
@@ -137,9 +137,8 @@ RSpec.describe DiscourseAi::Admin::AiToolsController do
 
   describe "#test" do
     it "runs an existing tool and returns the result" do
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+      post "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}/test.json",
            params: {
-             id: ai_tool.id,
              parameters: {
                input: "Hello, World!",
              },
@@ -149,43 +148,36 @@ RSpec.describe DiscourseAi::Admin::AiToolsController do
       expect(response.parsed_body["output"]).to eq("input" => "Hello, World!")
     end
 
-    it "runs a new unsaved tool and returns the result" do
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+    it "accept changes to the ai_tool parameters that redefine stuff" do
+      post "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}/test.json",
            params: {
              ai_tool: {
-               name: "New Tool",
-               description: "A new test tool",
-               script: "function invoke(params) { return 'New test result: ' + params.input; }",
-               parameters: [
-                 { name: "input", type: "string", description: "Input for the new test tool" },
-               ],
+               script: "function invoke(params) { return 'hi there'; }",
              },
-             parameters: {
-               input: "Test input",
-             },
-           }
-
-      expect(response.status).to eq(200)
-      expect(response.parsed_body["output"]).to eq("New test result: Test input")
-    end
-
-    it "returns an error for invalid tool_id" do
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
-           params: {
-             id: -1,
              parameters: {
                input: "Hello, World!",
              },
            }
 
-      expect(response.status).to eq(400)
-      expect(response.parsed_body["errors"]).to include("Couldn't find AiTool with 'id'=-1")
+      expect(response.status).to eq(200)
+      expect(response.parsed_body["output"]).to eq("hi there")
+    end
+
+    it "returns an error for invalid tool_id" do
+      post "/admin/plugins/discourse-ai/ai-tools/-1/test.json",
+           params: {
+             parameters: {
+               input: "Hello, World!",
+             },
+           }
+
+      expect(response.status).to eq(404)
     end
 
     it "handles exceptions during tool execution" do
       ai_tool.update!(script: "function invoke(params) { throw new Error('Test error'); }")
 
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+      post "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}/test.json",
            params: {
              id: ai_tool.id,
              parameters: {
diff --git a/spec/system/ai_bot/tool_spec.rb b/spec/system/ai_bot/tool_spec.rb
index 2d250345..ec52ca16 100644
--- a/spec/system/ai_bot/tool_spec.rb
+++ b/spec/system/ai_bot/tool_spec.rb
@@ -49,7 +49,8 @@ describe "AI Tool Management", type: :system do
     expect(page.first(".parameter-row__required-toggle").checked?).to eq(true)
     expect(page.first(".parameter-row__enum-toggle").checked?).to eq(false)
 
-    ensure_can_run_test
+    # not allowed to test yet
+    expect(page).not_to have_button(".ai-tool-editor__test-button")
 
     expect(page).not_to have_button(".ai-tool-editor__delete")
     find(".ai-tool-editor__save").click