diff --git a/tests/tools/test_ckg_tool.py b/tests/tools/test_ckg_tool.py new file mode 100644 index 000000000..2e0546c8d --- /dev/null +++ b/tests/tools/test_ckg_tool.py @@ -0,0 +1,73 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# SPDX-License-Identifier: MIT + +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import MagicMock, patch + +from trae_agent.tools.base import ToolCallArguments +from trae_agent.tools.ckg_tool import CKGTool + + +class TestCKGTool(unittest.IsolatedAsyncioTestCase): + async def test_execute_accepts_file_path_inside_codebase(self): + with TemporaryDirectory() as tmpdir: + code_file = Path(tmpdir) / "example.py" + code_file.write_text("def target():\n return 1\n") + + ckg_database = MagicMock() + ckg_database.query_function.return_value = [] + + with patch( + "trae_agent.tools.ckg_tool.CKGDatabase", return_value=ckg_database + ) as ckg_database_cls: + result = await CKGTool().execute( + ToolCallArguments( + { + "command": "search_function", + "path": str(code_file), + "identifier": "target", + "print_body": False, + } + ) + ) + + self.assertIsNone(result.error) + self.assertEqual(result.output, "No functions named target found.") + ckg_database_cls.assert_called_once_with(code_file.parent) + + async def test_execute_uses_git_root_for_file_path(self): + with TemporaryDirectory() as tmpdir: + codebase_root = Path(tmpdir) + nested_dir = codebase_root / "pkg" + nested_dir.mkdir() + code_file = nested_dir / "example.py" + code_file.write_text("def target():\n return 1\n") + + ckg_database = MagicMock() + ckg_database.query_function.return_value = [] + git_result = MagicMock(returncode=0, stdout=f"{codebase_root}\n") + + with ( + patch("trae_agent.tools.ckg_tool.subprocess.run", return_value=git_result), + patch( + "trae_agent.tools.ckg_tool.CKGDatabase", return_value=ckg_database + ) as ckg_database_cls, + ): + result = await CKGTool().execute( + ToolCallArguments( + { + "command": "search_function", + "path": str(code_file), + "identifier": "target", + } + ) + ) + + self.assertIsNone(result.error) + ckg_database_cls.assert_called_once_with(codebase_root) + + +if __name__ == "__main__": + unittest.main() diff --git a/trae_agent/tools/ckg_tool.py b/trae_agent/tools/ckg_tool.py index 49da80ddc..17a15fe67 100644 --- a/trae_agent/tools/ckg_tool.py +++ b/trae_agent/tools/ckg_tool.py @@ -1,6 +1,7 @@ # Copyright (c) 2025 ByteDance Ltd. and/or its affiliates # SPDX-License-Identifier: MIT +import subprocess from pathlib import Path from typing import override @@ -60,7 +61,7 @@ def get_parameters(self) -> list[ToolParameter]: ToolParameter( name="path", type="string", - description="The path to the codebase.", + description="The path to the codebase root or to a file inside the codebase.", required=True, ), ToolParameter( @@ -105,7 +106,9 @@ async def execute(self, arguments: ToolCallArguments) -> ToolExecResult: error=f"Codebase path {path} does not exist", error_code=-1, ) - if not codebase_path.is_dir(): + if codebase_path.is_file(): + codebase_path = self._resolve_codebase_path_from_file(codebase_path) + elif not codebase_path.is_dir(): return ToolExecResult( error=f"Codebase path {path} is not a directory", error_code=-1, @@ -132,6 +135,23 @@ async def execute(self, arguments: ToolCallArguments) -> ToolExecResult: case _: return ToolExecResult(error=f"Invalid command: {command}", error_code=-1) + def _resolve_codebase_path_from_file(self, file_path: Path) -> Path: + parent = file_path.parent + try: + result = subprocess.run( + ["git", "rev-parse", "--show-toplevel"], + cwd=parent, + capture_output=True, + text=True, + timeout=5, + check=False, + ) + if result.returncode == 0 and result.stdout.strip(): + return Path(result.stdout.strip()) + except (FileNotFoundError, subprocess.TimeoutExpired): + pass + return parent + def _search_function( self, ckg_database: CKGDatabase, identifier: str, print_body: bool = True ) -> str: