WIP - improvements to token usage tracking, simulator added for live testing, improvements to file loading

tests/test_conversation_memory.py

@@ -166,7 +166,7 @@ class TestConversationMemory:
             initial_context={},
         )
 
-        history = build_conversation_history(context)
+        history, tokens = build_conversation_history(context)
 
         # Test basic structure
         assert "CONVERSATION HISTORY" in history
@@ -207,8 +207,9 @@ class TestConversationMemory:
             initial_context={},
         )
 
-        history = build_conversation_history(context)
+        history, tokens = build_conversation_history(context)
         assert history == ""
+        assert tokens == 0
 
 
 class TestConversationFlow:
@@ -373,7 +374,7 @@ class TestConversationFlow:
                 initial_context={},
             )
 
-            history = build_conversation_history(context)
+            history, tokens = build_conversation_history(context)
             expected_turn_text = f"Turn {test_max}/{MAX_CONVERSATION_TURNS}"
             assert expected_turn_text in history
 
@@ -595,7 +596,7 @@ class TestConversationFlow:
             initial_context={"prompt": "Analyze this codebase", "files": ["/project/src/"]},
         )
 
-        history = build_conversation_history(final_context)
+        history, tokens = build_conversation_history(final_context)
 
         # Verify chronological order and speaker identification
         assert "--- Turn 1 (Gemini using analyze) ---" in history
@@ -670,7 +671,7 @@ class TestConversationFlow:
         mock_client.get.return_value = context_with_followup.model_dump_json()
 
         # Build history to verify follow-up is preserved
-        history = build_conversation_history(context_with_followup)
+        history, tokens = build_conversation_history(context_with_followup)
         assert "Found potential issue in authentication" in history
         assert "[Gemini's Follow-up: Should I examine the authentication middleware?]" in history
 
@@ -762,7 +763,7 @@ class TestConversationFlow:
             )
 
             # Build conversation history (should handle token limits gracefully)
-            history = build_conversation_history(context)
+            history, tokens = build_conversation_history(context)
 
             # Verify the history was built successfully
             assert "=== CONVERSATION HISTORY ===" in history

tests/test_cross_tool_continuation.py

@@ -247,7 +247,7 @@ class TestCrossToolContinuation:
         # Build conversation history
         from utils.conversation_memory import build_conversation_history
 
-        history = build_conversation_history(thread_context)
+        history, tokens = build_conversation_history(thread_context)
 
         # Verify tool names are included in the history
         assert "Turn 1 (Gemini using test_analysis)" in history

tests/test_large_prompt_handling.py

@@ -214,15 +214,15 @@ class TestLargePromptHandling:
             mock_model.generate_content.return_value = mock_response
             mock_create_model.return_value = mock_model
 
-            # Mock read_files to avoid file system access
-            with patch("tools.chat.read_files") as mock_read_files:
-                mock_read_files.return_value = "File content"
+            # Mock the centralized file preparation method to avoid file system access
+            with patch.object(tool, "_prepare_file_content_for_prompt") as mock_prepare_files:
+                mock_prepare_files.return_value = "File content"
 
                 await tool.execute({"prompt": "", "files": [temp_prompt_file, other_file]})
 
                 # Verify prompt.txt was removed from files list
-                mock_read_files.assert_called_once()
-                files_arg = mock_read_files.call_args[0][0]
+                mock_prepare_files.assert_called_once()
+                files_arg = mock_prepare_files.call_args[0][0]
                 assert len(files_arg) == 1
                 assert files_arg[0] == other_file
 

tests/test_precommit.py

@@ -228,10 +228,8 @@ class TestPrecommitTool:
     @patch("tools.precommit.find_git_repositories")
     @patch("tools.precommit.get_git_status")
     @patch("tools.precommit.run_git_command")
-    @patch("tools.precommit.read_files")
     async def test_files_parameter_with_context(
         self,
-        mock_read_files,
         mock_run_git,
         mock_status,
         mock_find_repos,
@@ -254,14 +252,15 @@ class TestPrecommitTool:
             (True, ""),  # unstaged files list (empty)
         ]
 
-        # Mock read_files
-        mock_read_files.return_value = "=== FILE: config.py ===\nCONFIG_VALUE = 42\n=== END FILE ==="
+        # Mock the centralized file preparation method  
+        with patch.object(tool, "_prepare_file_content_for_prompt") as mock_prepare_files:
+            mock_prepare_files.return_value = "=== FILE: config.py ===\nCONFIG_VALUE = 42\n=== END FILE ==="
 
-        request = PrecommitRequest(
-            path="/absolute/repo/path",
-            files=["/absolute/repo/path/config.py"],
-        )
-        result = await tool.prepare_prompt(request)
+            request = PrecommitRequest(
+                path="/absolute/repo/path",
+                files=["/absolute/repo/path/config.py"],
+            )
+            result = await tool.prepare_prompt(request)
 
         # Verify context files are included
         assert "## Context Files Summary" in result
@@ -316,9 +315,9 @@ class TestPrecommitTool:
             (True, ""),  # unstaged files (empty)
         ]
 
-        # Mock read_files to return empty (file not found)
-        with patch("tools.precommit.read_files") as mock_read:
-            mock_read.return_value = ""
+        # Mock the centralized file preparation method to return empty (file not found)
+        with patch.object(tool, "_prepare_file_content_for_prompt") as mock_prepare_files:
+            mock_prepare_files.return_value = ""
             result_with_files = await tool.prepare_prompt(request_with_files)
 
         assert "If you need additional context files" not in result_with_files

tests/test_prompt_regression.py

@@ -67,16 +67,16 @@ class TestPromptRegression:
             mock_model.generate_content.return_value = mock_model_response()
             mock_create_model.return_value = mock_model
 
-            # Mock file reading
-            with patch("tools.chat.read_files") as mock_read_files:
-                mock_read_files.return_value = "File content here"
+            # Mock file reading through the centralized method
+            with patch.object(tool, "_prepare_file_content_for_prompt") as mock_prepare_files:
+                mock_prepare_files.return_value = "File content here"
 
                 result = await tool.execute({"prompt": "Analyze this code", "files": ["/path/to/file.py"]})
 
                 assert len(result) == 1
                 output = json.loads(result[0].text)
                 assert output["status"] == "success"
-                mock_read_files.assert_called_once_with(["/path/to/file.py"])
+                mock_prepare_files.assert_called_once_with(["/path/to/file.py"], None, "Context files")
 
     @pytest.mark.asyncio
     async def test_thinkdeep_normal_analysis(self, mock_model_response):