refactor: rename flow enums and improve tool interaction logic

This commit is contained in:
2026-07-09 21:55:23 +08:00
parent 71a00056bf
commit 476a555f22
3 changed files with 86 additions and 91 deletions
@@ -1,20 +1,14 @@
use crate::agent::run::Flow;
use crate::agent::run::{TextFlow, ToolCallFlow};
use crate::agent::tool::ToolCall;
pub trait Hook: Send + Sync {
fn before_tool_call(&self, tool_call: &ToolCall) -> Flow;
/// Called after a tool call completes, before the JSON output is serialized to a string.
/// The `output` value can be modified in place.
/// Return `Flow::Stop` to abort the agent loop, `Flow::Skip` to discard this result,
/// or `Flow::Continue` to proceed normally.
fn after_tool_call(&self, _tool_name: &str, _output: &mut serde_json::Value) -> Flow {
Flow::Continue
}
}
impl Hook for () {
fn before_tool_call(&self, _tool_call: &ToolCall) -> Flow {
Flow::Continue
fn on_text_reply(&self, text: &str) -> TextFlow {
TextFlow::Continue
}
fn before_tool_call(&self, tool_call: &ToolCall) -> ToolCallFlow {
ToolCallFlow::Continue
}
fn after_tool_call(&self, tool_name: &str, output: &mut serde_json::Value) -> ToolCallFlow {
ToolCallFlow::Continue
}
}
@@ -2,24 +2,29 @@ use crate::agent::agent::Agent;
use crate::agent::agent::LlmApiAdaptor;
use crate::agent::error::PromptError;
use crate::agent::hook::Hook;
use crate::agent::memory::{AgentMemory, VecMemory};
use crate::agent::tool::{ToolCall, ToolDyn, ToolExecutionResult, ToolFailure, ToolOutcome};
use serde::Deserialize;
use serde_json::Value;
use std::pin::Pin;
use crate::agent::run::PromptOutput::{Text, Tool};
#[derive(Debug, Clone)]
pub struct PromptOutput {
pub text: String,
pub tool_calls: Vec<ToolCall>,
pub enum PromptOutput {
Text(String),
Tool(ToolCall),
}
/// Flow control for tool execution within the agent loop.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Flow {
pub enum TextFlow {
Continue,
Stop,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ToolCallFlow {
Continue,
/// Skip this tool call but continue processing other responses.
Skip,
/// Stop the agent execution entirely.
Stop,
}
@@ -69,45 +74,35 @@ where
M: LlmApiAdaptor<Message> + Send + Sync + 'a,
Message: Send + Sync + Clone + 'a,
{
type Output = Result<PromptOutput, PromptError>;
type Output = Result<Vec<PromptOutput>, PromptError>;
type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
fn into_future(self) -> Self::IntoFuture {
let agent = self.agent;
let message = self.message;
let per_prompt_system_prompt = self.system_prompt;
let extra_hooks = self.hooks;
let max_turns = self.max_turns;
Box::pin(async move {
let mut memory: Vec<Message> = Vec::new();
let mut memory = agent
.memory
.unwrap_or_else(|| VecMemory::new(Vec::new()).into());
let mut prompt_result: Vec<PromptOutput> = Vec::new();
// Pre-populate from agent's memory backend if available
if let Some(ref mem_backend) = agent.memory {
memory.extend(mem_backend.get_memory().iter().cloned());
}
let history = match agent.system_prompt {
Some(ref sp) => [sp].iter().chain(memory.iter().chain(message.into())),
None => memory.iter().chain(message.into()),
};
for _ in 0..max_turns {
// TODO perf issue for copy cost
let text = agent.model.complete(&history).await?;
// Effective system prompt: per-prompt override > agent default
match per_prompt_system_prompt {
Some(sp) => memory.push(sp),
None => {
if let Some(ref sp) = agent.system_prompt {
memory.push(sp.clone());
}
}
}
memory.push(message);
for turn in 0..max_turns {
// Delegate completion to the model (rig's CompletionModel pattern)
let text = agent.model.complete(&memory).await?;
// Try to parse the LLM reply as JSON (handle markdown fences)
// Try to parse the LLM reply as JSON (handle Markdown fences)
let cleaned = clean_json_response(&text);
match serde_json::from_str::<LlmJsonResponse>(&cleaned) {
Ok(json_resp) => {
// Check for <end/> marker — stop the turn immediately
// Check for <end/> marker — stop the _turn immediately
let reply_has_end = json_resp.reply_text.contains("<end/>");
let clean_text = json_resp
.reply_text
@@ -115,7 +110,19 @@ where
.trim()
.to_string();
// Run on_text_reply hooks
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
match hook.on_text_reply(&clean_text) {
TextFlow::Stop => {
return Err(PromptError::ToolError(
"text reply rejected by hook".to_string(),
));
}
TextFlow::Continue => {}
}
}
memory.push(agent.model.build_assistant_message(&clean_text));
prompt_result.push(Text(clean_text));
if reply_has_end {
return Ok(PromptOutput {
@@ -129,7 +136,7 @@ where
.into_iter()
.enumerate()
.map(|(idx, tc)| ToolCall {
id: format!("call_{turn}_{idx}"),
id: format!("{idx}"),
name: tc.tool_name,
args: tc.args,
})
@@ -143,20 +150,28 @@ where
}
for (tc_id, tc) in tool_calls.iter().enumerate() {
match run_hooks(&agent.hooks, &extra_hooks, tc) {
Flow::Stop => {
// inline run_hooks: before_tool_call hook
let mut should_skip = false;
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
match hook.before_tool_call(tc) {
ToolCallFlow::Stop => {
return Err(PromptError::ToolError(
"tool call rejected by hook".to_string(),
));
}
Flow::Skip => {
ToolCallFlow::Skip => {
let msg = agent
.model
.tool_result_message(&tc.name, "(skipped by hook)");
memory.push(msg);
continue;
should_skip = true;
break;
}
Flow::Continue => {}
ToolCallFlow::Continue => {}
}
}
if should_skip {
continue;
}
let result = {
@@ -181,21 +196,21 @@ where
};
match result.outcome {
ToolOutcome::Success => {
ToolOutcome::InternalOk => {
let mut json_output = result.output;
// Run after_tool_call hooks to allow output modification
for hook in agent.hooks.iter().chain(extra_hooks.iter()) {
match hook.after_tool_call(&tc.name, &mut json_output) {
Flow::Stop => {
ToolCallFlow::Stop => {
return Err(PromptError::ToolError(
"tool call output caused this turn stopped by hook".to_string(),
"tool call output caused this _turn stopped by hook".to_string(),
));
}
Flow::Skip => {
ToolCallFlow::Skip => {
json_output = serde_json::json!({"message":"tool call is ignored by hook"});
break;
}
Flow::Continue => {}
ToolCallFlow::Continue => {}
}
}
let arg_json = serde_json::to_string(&tc.args)
@@ -209,11 +224,12 @@ where
let msg =
agent.model.tool_result_message(&tc.name, &output_str);
memory.push(msg);
prompt_result.push(Tool(tc))
}
ToolOutcome::Failure(failure) if failure.fatal => {
return Err(PromptError::InternalError(failure.message));
ToolOutcome::InternalError(failure) if failure.fatal => {
return Err(PromptError::ToolError(failure.message));
}
ToolOutcome::Failure(failure) => {
ToolOutcome::InternalError(failure) => {
let msg = agent.model.tool_result_message(
&tc.name,
&format!("error: {}", failure.message),
@@ -224,7 +240,7 @@ where
}
}
Err(_) => {
// Not valid JSON — push as assistant message and continue to next turn
// Not valid JSON — push as assistant message and continue to next _turn
memory.push(agent.model.build_assistant_message(&text));
continue;
}
@@ -236,21 +252,6 @@ where
}
}
fn run_hooks(
agent_hooks: &[Box<dyn Hook>],
extra_hooks: &[Box<dyn Hook>],
tool_call: &ToolCall,
) -> Flow {
for hook in agent_hooks.iter().chain(extra_hooks.iter()) {
match hook.before_tool_call(tool_call) {
Flow::Stop => return Flow::Stop,
Flow::Skip => return Flow::Skip,
Flow::Continue => {}
}
}
Flow::Continue
}
#[derive(Deserialize)]
struct LlmJsonResponse {
reply_text: String,
@@ -262,7 +263,7 @@ struct LlmJsonResponse {
struct LlmToolCallRequest {
tool_name: String,
#[serde(default)]
args: serde_json::Value,
args: Value,
}
pub(crate) fn clean_json_response(text: &str) -> String {
@@ -84,8 +84,8 @@ impl ToolFailure {
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ToolOutcome {
Success,
Failure(ToolFailure),
InternalOk,
InternalError(ToolFailure),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -98,21 +98,21 @@ impl ToolExecutionResult {
pub fn success(output: serde_json::Value) -> Self {
Self {
output,
outcome: ToolOutcome::Success,
outcome: ToolOutcome::InternalOk,
}
}
pub fn failed(output: serde_json::Value, failure: ToolFailure) -> Self {
Self {
output,
outcome: ToolOutcome::Failure(failure),
outcome: ToolOutcome::InternalError(failure),
}
}
pub fn failure(&self) -> Option<&ToolFailure> {
match &self.outcome {
ToolOutcome::Success => None,
ToolOutcome::Failure(failure) => Some(failure),
ToolOutcome::InternalOk => None,
ToolOutcome::InternalError(failure) => Some(failure),
}
}