From 6facfda8ea5e78f7dd89bda3a376e7dee44cc582 Mon Sep 17 00:00:00 2001
From: Daniel Han <michaelhan2050@gmail.com>
Date: Sun, 31 May 2026 10:24:26 +0000
Subject: [PATCH] Studio: fix remaining broken tool-handler tests

Running the RAG suite surfaced three more test-only breakages (production
is correct):
- The scope tests patched `tool.__import__("core.rag.retrieval", ...)`, but a
  module has no `__import__` attribute (AttributeError before the stub ran).
  tool.py imports retrieval lazily, so patch core.rag.retrieval.retrieve_hybrid
  at the source instead.
- test_format_hits_produces_fenced_chunks asserted a stale `score="..."`
  attribute; the chunk tag emits chunk_index/tokens, not score.
- The execute_tool dispatch stub rejected the `mode` kwarg the handler passes.

Full RAG suite now: 101 passed, 5 skipped (server-gated model tests).
---
 tests/python/test_rag_tool_handler.py | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/tests/python/test_rag_tool_handler.py b/tests/python/test_rag_tool_handler.py
index 4d46f54004..c3944b5792 100644
--- a/tests/python/test_rag_tool_handler.py
+++ b/tests/python/test_rag_tool_handler.py
@@ -57,11 +57,7 @@ def test_kb_takes_precedence_over_thread():
         captured["scope"] = scope
         return []
 
-    with patch.object(
-        tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
-        "retrieve_hybrid",
-        _stub_retrieve,
-    ):
+    with patch("core.rag.retrieval.retrieve_hybrid", _stub_retrieve):
         result = tool.search_knowledge_base(
             query = "x",
             scope_kb_id = "kb-abc",
@@ -82,11 +78,7 @@ def test_thread_scope_when_only_thread_set():
         captured["scope"] = scope
         return []
 
-    with patch.object(
-        tool.__import__("core.rag.retrieval", fromlist = ["retrieve_hybrid"]),
-        "retrieve_hybrid",
-        _stub_retrieve,
-    ):
+    with patch("core.rag.retrieval.retrieve_hybrid", _stub_retrieve):
         tool.search_knowledge_base(
             query = "x",
             scope_thread_id = "thread-xyz",
@@ -122,11 +114,11 @@ def test_format_hits_produces_fenced_chunks():
         },
     ]
     result = _format_hits_for_llm(hits)
-    assert '<chunk id="1" source="alpha.pdf" page="3" score="0.780"' in result
+    assert '<chunk id="1" source="alpha.pdf" page="3" chunk_index="12" tokens="42">' in result
     assert 'chunk_index="12"' in result
     assert 'tokens="42"' in result
     assert "first body\n</chunk>" in result
-    assert '<chunk id="2" source="beta.md" score="0.610">' in result
+    assert '<chunk id="2" source="beta.md">' in result
     # Blank line between blocks so the model can scan them.
     assert "</chunk>\n\n<chunk" in result
 
@@ -208,6 +200,7 @@ def test_execute_tool_dispatches_to_search_knowledge_base():
         reranker_model = None,
         default_top_k = 5,
         min_score = 0.0,
+        **kwargs,
     ):
         called["query"] = query
         called["top_k"] = top_k