Trae Agent 控制核心逻辑详解
概述
Trae Agent 的控制逻辑采用状态机驱动的循环执行模型 ,通过多轮对话与 LLM 交互,结合工具调用完成复杂任务。核心控制流程位于 BaseAgent 类中。
核心架构
scss
┌─────────────────────────────────────────────────────────────────┐
│ Agent 控制架构 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ ┌─────────────────────────────────────────────────────────┐ │
│ │ execute_task() 主循环 │ │
│ │ while step_number <= max_steps: │ │
│ │ 1. 创建 AgentStep │ │
│ │ 2. _run_llm_step() - LLM 交互 │ │
│ │ 3. _finalize_step() - 步骤收尾 │ │
│ │ 4. 检查完成状态 │ │
│ └─────────────────────────────────────────────────────────┘ │
│ │ │
│ ┌───────────┼───────────┐ │
│ │ │ │ │
│ ▼ ▼ ▼ │
│ ┌────────────┐ ┌──────────┐ ┌──────────┐ │
│ │ 状态管理 │ │ 工具调用 │ │ 轨迹记录 │ │
│ │ │ │ │ │ │ │
│ │ AgentState │ │ ToolExec │ │ Trajectory│ │
│ │ AgentStep │ │ utor │ │ Recorder │ │
│ └────────────┘ └──────────┘ └──────────┘ │
│ │
└─────────────────────────────────────────────────────────────────┘
核心数据结构
1. AgentStep - 执行步骤
文件 : trae_agent/agent/agent_basics.py
python
@dataclass
class AgentStep:
"""表示 Agent 执行过程中的单个步骤"""
step_number: int # 步骤序号
state: AgentStepState # 步骤状态
thought: str | None = None # 思考内容
tool_calls: list[ToolCall] | None = None # 工具调用请求
tool_results: list[ToolResult] | None = None # 工具执行结果
llm_response: LLMResponse | None = None # LLM 响应
reflection: str | None = None # 反思内容
error: str | None = None # 错误信息
2. AgentExecution - 执行记录
python
@dataclass
class AgentExecution:
"""封装 Agent 任务的完整执行过程"""
task: str # 任务描述
steps: list[AgentStep] # 所有步骤
final_result: str | None = None # 最终结果
success: bool = False # 是否成功
total_tokens: LLMUsage | None = None # Token 使用量
execution_time: float = 0.0 # 执行时间
agent_state: AgentState = AgentState.IDLE # Agent 状态
3. 状态枚举
python
class AgentStepState(Enum):
"""步骤级别状态"""
THINKING = "thinking" # 思考中
CALLING_TOOL = "calling_tool" # 调用工具
REFLECTING = "reflecting" # 反思中
COMPLETED = "completed" # 已完成
ERROR = "error" # 出错
class AgentState(Enum):
"""Agent 级别状态"""
IDLE = "idle" # 空闲
RUNNING = "running" # 运行中
COMPLETED = "completed" # 已完成
ERROR = "error" # 出错
主控制流程
execute_task - 任务执行主循环
文件 : trae_agent/agent/base_agent.py
python
async def execute_task(self) -> AgentExecution:
"""执行任务的主入口"""
# 1. 初始化
start_time = time.time()
execution = AgentExecution(task=self._task, steps=[])
step: AgentStep | None = None
try:
messages = self._initial_messages
step_number = 1
execution.agent_state = AgentState.RUNNING
# 2. 主循环:最多执行 max_steps 步
while step_number <= self._max_steps:
# 创建新步骤
step = AgentStep(
step_number=step_number,
state=AgentStepState.THINKING
)
try:
# 3. 执行单步
messages = await self._run_llm_step(step, messages, execution)
# 4. 步骤收尾
await self._finalize_step(step, messages, execution)
# 5. 检查是否完成
if execution.agent_state == AgentState.COMPLETED:
break
step_number += 1
except Exception as error:
# 错误处理
execution.agent_state = AgentState.ERROR
step.state = AgentStepState.ERROR
step.error = str(error)
await self._finalize_step(step, messages, execution)
break
# 6. 检查是否超出最大步数
if step_number > self._max_steps and not execution.success:
execution.final_result = "Task execution exceeded maximum steps..."
execution.agent_state = AgentState.ERROR
except Exception as e:
execution.final_result = f"Agent execution failed: {str(e)}"
finally:
# 7. 清理资源
if self.docker_manager and not self.docker_keep:
self.docker_manager.stop()
await self._close_tools()
await self.cleanup_mcp_clients()
# 8. 计算执行时间并返回
execution.execution_time = time.time() - start_time
self._update_cli_console(step, execution)
return execution
单步执行流程
_run_llm_step - LLM 交互步骤
python
async def _run_llm_step(
self,
step: AgentStep,
messages: list[LLMMessage],
execution: AgentExecution
) -> list[LLMMessage]:
"""执行单轮 LLM 交互"""
# 1. 设置状态为 THINKING
step.state = AgentStepState.THINKING
self._update_cli_console(step, execution)
# 2. 调用 LLM
llm_response = self._llm_client.chat(
messages,
self._model_config,
self._tools
)
step.llm_response = llm_response
# 3. 更新显示
self._update_cli_console(step, execution)
# 4. 更新 Token 使用量
self._update_llm_usage(llm_response, execution)
# 5. 检查任务是否完成
if self.llm_indicates_task_completed(llm_response):
if self._is_task_completed(llm_response):
# 任务真正完成
execution.agent_state = AgentState.COMPLETED
execution.final_result = llm_response.content
execution.success = True
return messages
else:
# LLM 说完成但实际没完成
execution.agent_state = AgentState.RUNNING
return [LLMMessage(
role="user",
content=self.task_incomplete_message()
)]
else:
# 6. 有工具调用,处理工具调用
tool_calls = llm_response.tool_calls
return await self._tool_call_handler(tool_calls, step)
工具调用处理
_tool_call_handler - 工具调用处理器
python
async def _tool_call_handler(
self,
tool_calls: list[ToolCall] | None,
step: AgentStep
) -> list[LLMMessage]:
"""处理 LLM 返回的工具调用"""
messages: list[LLMMessage] = []
# 1. 检查是否有工具调用
if not tool_calls or len(tool_calls) <= 0:
messages = [LLMMessage(
role="user",
content="It seems that you have not completed the task.",
)]
return messages
# 2. 设置状态为 CALLING_TOOL
step.state = AgentStepState.CALLING_TOOL
step.tool_calls = tool_calls
self._update_cli_console(step)
# 3. 执行工具调用(并行或串行)
if self._model_config.parallel_tool_calls:
tool_results = await self._tool_caller.parallel_tool_call(tool_calls)
else:
tool_results = await self._tool_caller.sequential_tool_call(tool_calls)
step.tool_results = tool_results
self._update_cli_console(step)
# 4. 将工具结果添加到消息
for tool_result in tool_results:
message = LLMMessage(role="user", tool_result=tool_result)
messages.append(message)
# 5. 反思(如果工具执行失败)
reflection = self.reflect_on_result(tool_results)
if reflection:
step.state = AgentStepState.REFLECTING
step.reflection = reflection
self._update_cli_console(step)
messages.append(LLMMessage(role="assistant", content=reflection))
return messages
步骤状态流转
scss
┌─────────────┐
│ START │
└──────┬──────┘
│
▼
┌─────────────┐ LLM 调用 ┌─────────────┐
│ THINKING │ ────────────────▶ │ 有工具 │
│ (思考中) │ │ 调用? │
└─────────────┘ └──────┬──────┘
│ │
│ 无工具调用 │ 是
▼ ▼
┌─────────────┐ ┌─────────────┐
│ COMPLETED │ │ CALLING_TOOL│
│ (任务完成) │ │ (调用工具) │
└─────────────┘ └──────┬──────┘
│
▼
┌─────────────┐
│ 有失败? │
└──────┬──────┘
│
┌──────────────┼──────────────┐
│ │ │
│ 否 │ 是 │
▼ ▼ ▼
┌─────────────┐ ┌─────────────┐ ┌─────────────┐
│ 返回消息 │ │ REFLECTING │ │ ERROR │
│ 继续下轮 │ │ (反思中) │ │ (出错) │
└─────────────┘ └──────┬──────┘ └─────────────┘
│
▼
┌─────────────┐
│ 返回反思 │
│ 继续下轮 │
└─────────────┘
任务完成检测
llm_indicates_task_completed
python
def llm_indicates_task_completed(self, llm_response: LLMResponse) -> bool:
"""
检查 LLM 是否表示任务已完成
通过关键词匹配判断
"""
completion_indicators = [
"task completed",
"task finished",
"done",
"completed successfully",
"finished successfully",
]
response_lower = llm_response.content.lower()
return any(indicator in response_lower for indicator in completion_indicators)
_is_task_completed
python
def _is_task_completed(self, llm_response: LLMResponse) -> bool:
"""
验证任务是否真的完成
子类可以重写此方法添加自定义验证逻辑
"""
return True # 默认信任 LLM 的判断
步骤收尾处理
_finalize_step
python
async def _finalize_step(
self,
step: AgentStep,
messages: list[LLMMessage],
execution: AgentExecution
) -> None:
"""步骤收尾处理"""
# 1. 设置步骤状态为完成
step.state = AgentStepState.COMPLETED
# 2. 记录轨迹
self._record_handler(step, messages)
# 3. 更新 CLI 显示
self._update_cli_console(step, execution)
# 4. 添加到执行记录
execution.steps.append(step)
轨迹记录
_record_handler
python
def _record_handler(self, step: AgentStep, messages: list[LLMMessage]) -> None:
"""记录执行轨迹"""
if self.trajectory_recorder:
self.trajectory_recorder.record_agent_step(
step_number=step.step_number,
state=step.state.value,
llm_messages=messages,
llm_response=step.llm_response,
tool_calls=step.tool_calls,
tool_results=step.tool_results,
reflection=step.reflection,
error=step.error,
)
完整执行流程示例
css
用户任务: "修复 src/main.py 中的 bug"
初始化:
messages = [
{role: "system", content: "You are..."},
{role: "user", content: "修复 src/main.py 中的 bug"}
]
第 1 步:
├─ 状态: THINKING
├─ LLM 调用: "我来查看代码..."
├─ 工具调用: view({"path": "/repo/src/main.py"})
├─ 状态: CALLING_TOOL
├─ 工具结果: 文件内容...
├─ 状态: COMPLETED
└─ 返回 messages (包含工具结果)
第 2 步:
├─ 状态: THINKING
├─ LLM 调用: "我发现了问题..."
├─ 工具调用: str_replace({"path": "/repo/src/main.py", ...})
├─ 状态: CALLING_TOOL
├─ 工具结果: 替换成功
├─ 状态: COMPLETED
└─ 返回 messages
第 3 步:
├─ 状态: THINKING
├─ LLM 调用: "让我验证修复..."
├─ 工具调用: bash({"command": "python test.py"})
├─ 状态: CALLING_TOOL
├─ 工具结果: 失败 (exit code 1)
├─ 状态: REFLECTING
├─ 反思: "The tool execution failed..."
├─ 状态: COMPLETED
└─ 返回 messages (包含反思)
第 4 步:
├─ 状态: THINKING
├─ LLM 调用: "测试失败了,我需要..."
├─ 工具调用: str_replace({"path": "/repo/src/main.py", ...})
├─ ...
第 N 步:
├─ 状态: THINKING
├─ LLM 调用: "Task completed! The bug..."
├─ 检测完成: llm_indicates_task_completed() = True
├─ 验证完成: _is_task_completed() = True
├─ 状态: COMPLETED
└─ 执行结束
结果:
execution.success = True
execution.final_result = "Task completed! The bug..."
execution.steps = [Step1, Step2, Step3, Step4, ..., StepN]
核心设计原则
| 原则 | 说明 |
|---|---|
| 状态驱动 | 每个步骤都有明确的状态,便于跟踪和调试 |
| 循环执行 | 通过 while 循环实现多轮对话,直到任务完成 |
| 工具驱动 | 以工具调用为核心交互方式,LLM 通过工具影响环境 |
| 错误恢复 | 通过反思机制自动处理工具执行失败 |
| 可扩展 | 关键方法可重写,支持自定义逻辑 |
核心文件索引
| 文件 | 职责 | 重要性 |
|---|---|---|
agent/base_agent.py |
Agent 控制核心逻辑 | ⭐⭐⭐ |
agent/agent_basics.py |
基础数据结构和状态定义 | ⭐⭐⭐ |
agent/trae_agent.py |
TraeAgent 具体实现 | ⭐⭐⭐ |
最后更新: 2026-03-16