Trae Agent 源码阅读

trae-agent/
├── trae_agent/                       # 核心包目录
│   ├── agent/                        # Agent核心逻辑
│   │   ├── agent_basics.py           # 基础数据结构 (AgentState, AgentStep)
│   │   ├── base_agent.py             # 抽象基类 (ReAct循环逻辑)
│   │   ├── trae_agent.py             # 具体实现 (软件工程专用)
│   │   ├── agent.py                  # 包装器 (外观模式)
│   │   └── docker_manager.py         # Docker环境管理
│   │
│   ├── tools/                        # 工具集合
│   │   ├── base.py                   # 工具基类和接口
│   │   ├── bash_tool.py              # Shell命令执行
│   │   ├── edit_tool.py              # 文件编辑工具
│   │   ├── json_edit_tool.py         # JSON编辑工具
│   │   ├── task_done_tool.py         # 任务完成工具
│   │   ├── sequential_thinking_tool.py # 结构化思考
│   │   ├── mcp_tool.py               # MCP协议工具
│   │   └── ckg/                      # 代码知识图谱
│   │
│   ├── utils/                        # 工具类模块
│   │   ├── llm_clients/              # LLM客户端
│   │   │   ├── llm_client.py         # 统一客户端接口
│   │   │   ├── anthropic_client.py   # Claude集成
│   │   │   ├── openai_client.py      # GPT集成
│   │   │   ├── google_client.py      # Gemini集成
│   │   │   └── ...                   # 其他LLM提供商
│   │   │
│   │   ├── cli/                      # 命令行界面
│   │   │   ├── console_factory.py    # 控制台工厂
│   │   │   ├── simple_console.py     # 简单文本界面
│   │   │   └── rich_console.py       # 富文本界面
│   │   │
│   │   ├── config.py                 # 配置管理
│   │   ├── trajectory_recorder.py    # 轨迹记录
│   │   └── mcp_client.py             # MCP协议客户端
│   │
│   ├── prompt/                       # 提示词管理
│   │   └── agent_prompt.py           # Agent提示词
│   │
│   ├── cli.py                        # 命令行入口
│   └── dist/                         # 分发工具
│
├── evaluation/                       # 评估框架
│   ├── run_evaluation.py             # SWE-bench评估脚本
│   ├── utils.py                      # 评估工具函数
│   └── patch_selection/              # 补丁选择评估
│
├── tests/                            # 测试套件
│   ├── agent/                        # Agent测试
│   ├── tools/                        # 工具测试
│   └── utils/                        # 工具类测试
│
├── docs/                             # 文档
│   ├── tools.md                      # 工具文档
│   └── TRAJECTORY_RECORDING.md       # 轨迹记录文档
│
├── server/                           # FastAPI服务端 (可选)
│
├── pyproject.toml                    # Python项目配置
├── trae_config.yaml.example          # 配置模板
└── README.md                         # 项目说明

@override  
def _is_task_completed(self, llm_response: LLMResponse) -> bool:  
    """Enhanced task completion detection."""  
    if self.must_patch == "true":  
        model_patch = self.get_git_diff()  
        patch = self.remove_patches_to_tests(model_patch)  
        if not patch.strip():  
            return False  
  
    return True

def remove_patches_to_tests(self, model_patch: str) -> str:  
    """  
    Remove any changes to the tests directory from the provided patch.    This is to ensure that the model_patch does not disturb the repo's    tests when doing acceptance testing with the `test_patch`.    """    lines = model_patch.splitlines(keepends=True)  
    filtered_lines: list[str] = []  
    test_patterns = ["/test/", "/tests/", "/testing/", "test_", "tox.ini"]  
    is_tests = False  
  
    for line in lines:  
        if line.startswith("diff --git a/"):  
            target_path = line.split()[-1]  
            is_tests = target_path.startswith("b/") and any(  
                p in target_path for p in test_patterns  
            )  
  
        if not is_tests:  
            filtered_lines.append(line)  
  
    return "".join(filtered_lines)

class AgentState(Enum):  
    """Defines possible states during an agent's execution lifecycle."""  
  
    IDLE = "idle"  
    RUNNING = "running"  
    COMPLETED = "completed"  
    ERROR = "error"

class AgentStepState(Enum):  
    """Defines possible states during an agent's execution lifecycle."""  
  
    THINKING = "thinking"  
    CALLING_TOOL = "calling_tool"  
    REFLECTING = "reflecting"  
    COMPLETED = "completed"  
    ERROR = "error"  

class Tool(ABC):
    """Base class for all tools."""
	...

    @abstractmethod
    def get_name(self) -> str:
        """Get the tool name."""
        pass

    @abstractmethod
    def get_description(self) -> str:
        """Get the tool description."""
        pass

    @abstractmethod
    def get_parameters(self) -> list[ToolParameter]:
        """Get the tool parameters."""
        pass

    @abstractmethod
    async def execute(self, arguments: ToolCallArguments) -> ToolExecResult:
        """Execute the tool with given parameters."""
        pass
	...
	
@dataclass
class ToolParameter:
    """Tool parameter definition."""

    name: str
    type: str | list[str]
    description: str
    enum: list[str] | None = None
    items: dict[str, object] | None = None
    required: bool = True

-- 函数表
CREATE TABLE functions (
	id INTEGER PRIMARY KEY,
	NAME TEXT NOT NULL, -- 函数名
	file_path TEXT NOT NULL, -- 文件路径
	body TEXT NOT NULL, -- 函数体代码
	start_line INTEGER, -- 起始行号
	end_line INTEGER, -- 结束行号
	parent_function TEXT, -- 父函数（嵌套函数）
	parent_class TEXT -- 所属类（方法）
);

-- 类表
CREATE TABLE classes (
	id INTEGER PRIMARY KEY,
	NAME TEXT NOT NULL, -- 类名
	file_path TEXT NOT NULL, -- 文件路径
	body TEXT NOT NULL, -- 类体代码
	fields TEXT, -- 字段列表
	methods TEXT, -- 方法列表
	start_line INTEGER, -- 起始行号
	end_line INTEGER -- 结束行号
)

_sentinel: str = ",,,,bash-command-exit-__ERROR_CODE__-banner,,,," # `__ERROR_CODE__` will be replaced by `$?` or `!errorlevel!` later

sentinel_before, pivot, sentinel_after = self._sentinel.partition("__ERROR_CODE__")  
assert pivot == "__ERROR_CODE__"  
  
errcode_retriever = "!errorlevel!" if os.name == "nt" else "$?"  
command_sep = "&" if os.name == "nt" else ";"

# send command to the process  
self._process.stdin.write(  
    b"(\n"  
    + command.encode()  
    + f"\n){command_sep} echo {self._sentinel.replace('__ERROR_CODE__', errcode_retriever)}\n".encode()  
)  
await self._process.stdin.drain()  
  
# read output from the process, until the sentinel is found  
try:  
    async with asyncio.timeout(self._timeout):  
        while True:  
            await asyncio.sleep(self._output_delay)  
            # if we read directly from stdout/stderr, it will wait forever for  
            # EOF. use the StreamReader buffer directly instead.            output: str = self._process.stdout._buffer.decode()  # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType]  
            if sentinel_before in output:  
                # strip the sentinel from output  
                output, pivot, exit_banner = output.rpartition(sentinel_before)  
                assert pivot  
  
                # get error code inside banner  
                error_code_str, pivot, _ = exit_banner.partition(sentinel_after)  
                if not pivot or not error_code_str.isdecimal():  
                    continue  
  
                error_code = int(error_code_str)  
                break  
except asyncio.TimeoutError:  
    self._timed_out = True  
    raise ToolError(  
        f"timed out: bash has not returned in {self._timeout} seconds and must be restarted",  
    ) from None

def validate_path(self, command: str, path: Path):
	"""Validate the path for the str_replace_editor tool."""
	if not path.is_absolute():
		suggested_path = Path("/") / path
		raise ToolError(
			f"The path {path} is not an absolute path, it should start with `/`. Maybe you meant {suggested_path}?"
		)
	# Check if path exists
	if not path.exists() and command != "create":
		raise ToolError(f"The path {path} does not exist. Please provide a valid path.")
	if path.exists() and command == "create":
		raise ToolError(
			f"File already exists at: {path}. Cannot overwrite files using command `create`."
		)
	# Check if the path points to a directory
	if path.is_dir() and command != "view":
		raise ToolError(
			f"The path {path} is a directory and only the `view` command can be used on directories"
		)

def str_replace(self, path: Path, old_str: str, new_str: str | None) -> ToolExecResult:
	... # 文件内容替换

	# Create a snippet of the edited section
	replacement_line = file_content.split(old_str)[0].count("\n")
	start_line = max(0, replacement_line - SNIPPET_LINES)
	end_line = replacement_line + SNIPPET_LINES + new_str.count("\n")
	snippet = "\n".join(new_file_content.split("\n")[start_line : end_line + 1])

	# Prepare the success message
	success_msg = f"The file {path} has been edited. "
	success_msg += self._make_output(snippet, f"a snippet of {path}", start_line + 1)
	success_msg += "Review the changes and make sure they are as expected. Edit the file again if necessary."

	return ToolExecResult(
		output=success_msg,
	)

def _make_output(
	self,
	file_content: str,
	file_descriptor: str,
	init_line: int = 1,
	expand_tabs: bool = True,
):
	"""Generate output for the CLI based on the content of a file."""
	file_content = maybe_truncate(file_content)
	if expand_tabs:
		file_content = file_content.expandtabs()
	file_content = "\n".join(
		[f"{i + init_line:6}\t{line}" for i, line in enumerate(file_content.split("\n"))]
	)
	return (
		f"Here's the result of running `cat -n` on {file_descriptor}:\n" + file_content + "\n"
	)

A detailed tool for dynamic and reflective problem-solving through thoughts.
This tool helps analyze problems through a flexible thinking process that can adapt and evolve.
Each thought can build on, question, or revise previous insights as understanding deepens.

When to use this tool:
- Breaking down complex problems into steps
- Planning and design with room for revision
- Analysis that might need course correction
- Problems where the full scope might not be clear initially
- Problems that require a multi-step solution
- Tasks that need to maintain context over multiple steps
- Situations where irrelevant information needs to be filtered out

Key features:
- You can adjust total_thoughts up or down as you progress
- You can question or revise previous thoughts
- You can add more thoughts even after reaching what seemed like the end
- You can express uncertainty and explore alternative approaches
- Not every thought needs to build linearly - you can branch or backtrack
- Generates a solution hypothesis
- Verifies the hypothesis based on the Chain of Thought steps
- Repeats the process until satisfied
- Provides a correct answer

Parameters explained:
- thought: Your current thinking step, which can include:
* Regular analytical steps
* Revisions of previous thoughts
* Questions about previous decisions
* Realizations about needing more analysis
* Changes in approach
* Hypothesis generation
* Hypothesis verification
- next_thought_needed: True if you need more thinking, even if at what seemed like the end
- thought_number: Current number in sequence (can go beyond initial total if needed)
- total_thoughts: Current estimate of thoughts needed (can be adjusted up/down)
- is_revision: A boolean indicating if this thought revises previous thinking
- revises_thought: If is_revision is true, which thought number is being reconsidered
- branch_from_thought: If branching, which thought number is the branching point
- branch_id: Identifier for the current branch (if any)
- needs_more_thoughts: If reaching end but realizing more thoughts needed

You should:
1. Start with an initial estimate of needed thoughts, but be ready to adjust
2. Feel free to question or revise previous thoughts
3. Don't hesitate to add more thoughts if needed, even at the "end"
4. Express uncertainty when present
5. Mark thoughts that revise previous thinking or branch into new paths
6. Ignore information that is irrelevant to the current step
7. Generate a solution hypothesis when appropriate
8. Verify the hypothesis based on the Chain of Thought steps
9. Repeat the process until satisfied with the solution
10. Provide a single, ideally correct answer as the final output
11. Only set next_thought_needed to false when truly done and a satisfactory answer is reached

@dataclass
class ThoughtData:
	thought: str
	thought_number: int
	total_thoughts: int
	next_thought_needed: bool
	is_revision: bool | None = None
	revises_thought: int | None = None
	branch_from_thought: int | None = None
	branch_id: str | None = None
	needs_more_thoughts: bool | None = None

# ROLE: Act as an expert code evaluator. Given a codebase, an github issue and **{candidate_length} candidate patches** proposed by your colleagues, your responsibility is to **select the correct one** to solve the issue.

# WORK PROCESS:
You are given a software issue and multiple candidate patches. Your goal is to identify the patch that correctly resolves the issue.

Follow these steps methodically:

**1. Understand the Issue and Codebase**
Carefully read the issue description to comprehend the problem. You may need to examine the codebase for context, including:
    (1) Code referenced in the issue description;
    (2) The original code modified by each patch;
    (3) Unchanged parts of the same file;
    (4) Related files, functions, or modules that interact with the affected code.

**2. Analyze the Candidate Patches**
For each patch, analyze its logic and intended fix. Consider whether the changes align with the issue description and coding conventions.

**3. Validate Functionality (Optional but Recommended)**
If needed, write and run unit tests to evaluate the correctness and potential side effects of each patch.

**4. Select the Best Patch**
Choose the patch that best resolves the issue with minimal risk of introducing new problems.

# FINAL REPORT: If you have successfully selected the correct patch, submit your answer in the following format:
### Status: succeed
### Result: Patch-x
### Analysis: [Explain why Patch-x is correct.]

# IMPORTANT TIPS:
1. Never avoid making a selection.
2. Do not propose new patches.
3. There must be at least one correct patch.

# majority voting
if majority_voting:
	final_id_list, final_patch_list = [], []
	for idx in range(num_candidate):
		select_agent = SelectorAgent(
			llm_config=llm_config,
			sandbox=sandbox,
			project_path=project_path,
			issue_description=instance["problem_statement"],
			trajectory_file_name=get_trajectory_filename(
				instance["instance_id"], log_path, group_id, idx
			),
			candidate_list=candidate_list,
			max_turn=max_turn,
		)

		final_id, final_patch = select_agent.run()
		final_id_list.append(final_id)
		final_patch_list.append(final_patch)
		if max(Counter(final_id_list).values()) > num_candidate / 2:
			break
	print(f"[Retry No:{current_try}] majority voting done")
	sys.stdout.flush()
	sys.stderr.flush()

	counter = Counter(final_id_list)
	max_count = max(counter.values())
	most_common_ids = [
		elem for elem, count in counter.items() if count == max_count
	]
	result = {}
	for id_ in most_common_ids:
		indexes = [i for i, val in enumerate(final_id_list) if val == id_]
		result[id_] = indexes
	final_id = most_common_ids[0]
	final_patch = final_patch_list[result[final_id][0]]
	print(f"[Retry No:{current_try}] final_id_list: {final_id_list}")
	sys.stdout.flush()
	sys.stderr.flush()