1683 lines
61 KiB
Rust
1683 lines
61 KiB
Rust
use crate::agent::context_compressor::estimate_tokens;
|
||
use crate::agent::media_handler::MediaHandlerRegistry;
|
||
use crate::agent::system_prompt::build_system_prompt;
|
||
use crate::agent::turn_event::{AgentTurnContext, TurnEvent};
|
||
use crate::bus::message::ContentBlock;
|
||
use crate::bus::{ChatMessage, MediaRef};
|
||
use crate::config::LLMProviderConfig;
|
||
use crate::observability::{Observer, ObserverEvent, ToolExecutionOutcome, truncate_args};
|
||
use crate::providers::{
|
||
ChatCompletionRequest, ChatCompletionResponse, LLMProvider, Message, ProviderChunk,
|
||
ProviderResponseAccumulator, ToolCall, create_provider,
|
||
};
|
||
use crate::tools::{ToolExecutionContext, ToolRegistry};
|
||
use std::collections::VecDeque;
|
||
use std::hash::{Hash, Hasher};
|
||
use std::path::PathBuf;
|
||
use std::sync::Arc;
|
||
use std::time::Instant;
|
||
|
||
use futures_util::StreamExt;
|
||
|
||
/// Maximum characters in a tool result before truncation.
|
||
/// Prevents context overflow from large tool outputs.
|
||
const MAX_TOOL_RESULT_CHARS: usize = 16_000;
|
||
/// Minimum characters to keep when truncating
|
||
const TRUNCATION_SUFFIX_LEN: usize = 200;
|
||
const TOOL_PREVIEW_CHARS: usize = 1_000;
|
||
|
||
enum MediaOrigin<'a> {
|
||
User,
|
||
Tool(&'a str),
|
||
Message,
|
||
}
|
||
|
||
fn should_include_message_media(messages: &[ChatMessage], index: usize) -> bool {
|
||
let message = &messages[index];
|
||
if message.role != "tool" {
|
||
return true;
|
||
}
|
||
let active_tool_start = messages
|
||
.iter()
|
||
.rposition(|candidate| candidate.role != "tool")
|
||
.map_or(0, |last_non_tool| last_non_tool + 1);
|
||
index >= active_tool_start
|
||
}
|
||
|
||
/// Build content blocks from text and media, respecting model input capabilities
|
||
fn build_content_blocks(
|
||
text: &str,
|
||
media_refs: &[MediaRef],
|
||
input_types: &[String],
|
||
registry: &MediaHandlerRegistry,
|
||
origin: MediaOrigin<'_>,
|
||
) -> Vec<ContentBlock> {
|
||
let mut blocks = Vec::new();
|
||
|
||
if !media_refs.is_empty() {
|
||
let attachments = media_refs
|
||
.iter()
|
||
.map(|media_ref| {
|
||
let path = std::path::Path::new(&media_ref.path);
|
||
let name = path
|
||
.file_name()
|
||
.map(|name| name.to_string_lossy().into_owned())
|
||
.unwrap_or_else(|| media_ref.path.clone());
|
||
let extension = path
|
||
.extension()
|
||
.map(|extension| extension.to_string_lossy().into_owned());
|
||
let mime_type = mime_guess::from_path(path)
|
||
.first_or_octet_stream()
|
||
.to_string();
|
||
let size_bytes = std::fs::metadata(path).ok().map(|metadata| metadata.len());
|
||
let native_input = input_types.contains(&media_ref.media_type)
|
||
&& registry.supports(&media_ref.media_type);
|
||
serde_json::json!({
|
||
"name": name,
|
||
"extension": extension,
|
||
"media_type": media_ref.media_type,
|
||
"mime_type": mime_type,
|
||
"size_bytes": size_bytes,
|
||
"path": media_ref.path,
|
||
"content_delivery": if native_input {
|
||
"also included as a model-native content block"
|
||
} else {
|
||
"content is not embedded in this model request; path remains available to file tools"
|
||
},
|
||
})
|
||
})
|
||
.collect::<Vec<_>>();
|
||
let manifest = serde_json::Value::Array(attachments).to_string();
|
||
let note = match origin {
|
||
MediaOrigin::User if text.is_empty() => {
|
||
"用户发送了以下附件。path 是 Gateway 内部存储路径，可供文件工具读取；附件内容未必已嵌入模型输入。".to_string()
|
||
}
|
||
MediaOrigin::User => {
|
||
"随本条用户消息同时提交的附件。path 是 Gateway 内部存储路径，可供文件工具读取；content_delivery 说明附件内容是否另以模型原生内容块提供。".to_string()
|
||
}
|
||
MediaOrigin::Tool(tool_name) => format!(
|
||
"工具 {tool_name} 返回了以下附件。content_delivery 说明附件内容是否另以模型原生内容块提供。"
|
||
),
|
||
MediaOrigin::Message => {
|
||
"本条消息包含以下附件。content_delivery 说明附件内容是否另以模型原生内容块提供。".to_string()
|
||
}
|
||
};
|
||
let message_text = if text.is_empty() {
|
||
format!("[{note}]\n{manifest}")
|
||
} else {
|
||
format!("{text}\n\n[{note}]\n{manifest}")
|
||
};
|
||
blocks.push(ContentBlock::text(message_text));
|
||
|
||
for mr in media_refs {
|
||
if input_types.contains(&mr.media_type) && registry.supports(&mr.media_type) {
|
||
match registry.handle(&mr.media_type, &mr.path) {
|
||
Ok(content_blocks) => blocks.extend(content_blocks),
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
path = %mr.path,
|
||
media_type = %mr.media_type,
|
||
error = %e,
|
||
"Media handler failed, falling back to text placeholder"
|
||
);
|
||
blocks.push(ContentBlock::text(format!(
|
||
"[用户发来了一个文件，但处理失败: {}, 错误: {}]",
|
||
mr.path, e
|
||
)));
|
||
}
|
||
}
|
||
} else {
|
||
tracing::debug!(
|
||
path = %mr.path,
|
||
media_type = %mr.media_type,
|
||
model_input_types = ?input_types,
|
||
"Media type not supported by model; attachment manifest remains available"
|
||
);
|
||
}
|
||
}
|
||
} else if !text.is_empty() {
|
||
blocks.push(ContentBlock::text(text));
|
||
}
|
||
|
||
if blocks.is_empty() {
|
||
blocks.push(ContentBlock::text(""));
|
||
}
|
||
|
||
blocks
|
||
}
|
||
|
||
/// Truncate tool result if it exceeds MAX_TOOL_RESULT_CHARS.
|
||
/// Preserves the end of the output as it often contains the conclusion/useful result.
|
||
fn truncate_tool_result(output: &str) -> String {
|
||
if output.len() <= MAX_TOOL_RESULT_CHARS {
|
||
return output.to_string();
|
||
}
|
||
|
||
let truncated_start_len = output.len().saturating_sub(TRUNCATION_SUFFIX_LEN);
|
||
if truncated_start_len > MAX_TOOL_RESULT_CHARS {
|
||
// Even after removing suffix, still too long - take from beginning
|
||
format!(
|
||
"{}...\n\n[Output truncated - {} characters removed]",
|
||
&output[..output.ceil_char_boundary(MAX_TOOL_RESULT_CHARS - 100)],
|
||
output.len() - MAX_TOOL_RESULT_CHARS + 100
|
||
)
|
||
} else {
|
||
// Keep most of the end which usually contains the useful result
|
||
format!(
|
||
"...\n\n[Output truncated - {} characters removed]\n\n{}",
|
||
truncated_start_len,
|
||
&output[output.floor_char_boundary(truncated_start_len)..]
|
||
)
|
||
}
|
||
}
|
||
|
||
fn tool_result_preview(output: &str) -> String {
|
||
if output.len() <= TOOL_PREVIEW_CHARS {
|
||
output.to_string()
|
||
} else {
|
||
format!(
|
||
"{}…",
|
||
&output[..output.floor_char_boundary(TOOL_PREVIEW_CHARS)]
|
||
)
|
||
}
|
||
}
|
||
|
||
/// Loop detection result.
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
enum LoopDetectionResult {
|
||
/// No warning needed.
|
||
Ok,
|
||
/// Warning: same tool + args repeated N times.
|
||
Warning(String),
|
||
}
|
||
|
||
/// Configuration for loop detector.
|
||
#[derive(Debug, Clone)]
|
||
struct LoopDetectorConfig {
|
||
/// Master switch.
|
||
enabled: bool,
|
||
/// Warn every N consecutive identical calls.
|
||
warn_every: usize,
|
||
}
|
||
|
||
impl Default for LoopDetectorConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
enabled: true,
|
||
warn_every: 5,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// A single recorded tool invocation in the sliding window.
|
||
#[derive(Debug, Clone)]
|
||
struct ToolCallRecord {
|
||
name: String,
|
||
args_hash: u64,
|
||
}
|
||
|
||
/// Stateful loop detector that monitors for repetitive patterns.
|
||
struct LoopDetector {
|
||
config: LoopDetectorConfig,
|
||
window: VecDeque<ToolCallRecord>,
|
||
}
|
||
|
||
impl LoopDetector {
|
||
fn new(config: LoopDetectorConfig) -> Self {
|
||
Self {
|
||
window: VecDeque::with_capacity(config.warn_every * 2),
|
||
config,
|
||
}
|
||
}
|
||
|
||
/// Record a completed tool call and check for loop patterns.
|
||
/// Returns Warning every `warn_every` consecutive identical calls.
|
||
fn record(&mut self, name: &str, args: &serde_json::Value) -> LoopDetectionResult {
|
||
if !self.config.enabled {
|
||
return LoopDetectionResult::Ok;
|
||
}
|
||
|
||
let record = ToolCallRecord {
|
||
name: name.to_string(),
|
||
args_hash: hash_json_value(args),
|
||
};
|
||
|
||
// Maintain sliding window
|
||
if self.window.len() >= self.config.warn_every * 2 {
|
||
self.window.pop_front();
|
||
}
|
||
self.window.push_back(record);
|
||
|
||
// Count consecutive identical calls
|
||
let last = self.window.back().unwrap();
|
||
let consecutive: usize = self
|
||
.window
|
||
.iter()
|
||
.rev()
|
||
.take_while(|r| r.name == last.name && r.args_hash == last.args_hash)
|
||
.count();
|
||
|
||
// Warn every warn_every times
|
||
if consecutive > 0 && consecutive.is_multiple_of(self.config.warn_every) {
|
||
LoopDetectionResult::Warning(format!(
|
||
"注意: 工具 '{}' 已连续执行 {} 次，参数相同。如果任务没有进展，请尝试其他方法。",
|
||
last.name, consecutive
|
||
))
|
||
} else {
|
||
LoopDetectionResult::Ok
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Hash a JSON value deterministically (key-order independent).
|
||
fn hash_json_value(value: &serde_json::Value) -> u64 {
|
||
let mut hasher = std::collections::hash_map::DefaultHasher::new();
|
||
let canonical = canonicalise_json(value);
|
||
canonical.hash(&mut hasher);
|
||
hasher.finish()
|
||
}
|
||
|
||
/// Return a clone of value with all object keys sorted recursively.
|
||
fn canonicalise_json(value: &serde_json::Value) -> serde_json::Value {
|
||
match value {
|
||
serde_json::Value::Object(map) => {
|
||
let mut sorted: Vec<(&String, &serde_json::Value)> = map.iter().collect();
|
||
sorted.sort_by_key(|(k, _)| *k);
|
||
let new_map: serde_json::Map<String, serde_json::Value> = sorted
|
||
.into_iter()
|
||
.map(|(k, v)| (k.clone(), canonicalise_json(v)))
|
||
.collect();
|
||
serde_json::Value::Object(new_map)
|
||
}
|
||
serde_json::Value::Array(arr) => {
|
||
serde_json::Value::Array(arr.iter().map(canonicalise_json).collect())
|
||
}
|
||
other => other.clone(),
|
||
}
|
||
}
|
||
/// AgentLoop - Stateless agent that processes messages with tool calling support.
|
||
/// History is managed externally by SessionManager.
|
||
pub struct AgentLoop {
|
||
provider: Arc<dyn LLMProvider>,
|
||
tools: Arc<ToolRegistry>,
|
||
observer: Option<Arc<dyn Observer>>,
|
||
max_iterations: usize,
|
||
workspace_dir: PathBuf,
|
||
model_name: String,
|
||
context_window: usize,
|
||
input_types: Vec<String>,
|
||
media_registry: MediaHandlerRegistry,
|
||
}
|
||
#[derive(Debug, Clone)]
|
||
pub struct AgentProcessResult {
|
||
pub final_response: ChatMessage,
|
||
pub emitted_messages: Vec<ChatMessage>,
|
||
pub total_tokens: Option<u32>,
|
||
pub usage: Option<crate::providers::Usage>,
|
||
/// Provider usage for the final successful request in this Turn. This is
|
||
/// the correct basis for context-window occupancy; `usage` is accumulated
|
||
/// across every tool iteration.
|
||
pub last_request_usage: Option<crate::providers::Usage>,
|
||
}
|
||
|
||
fn merge_usage(total: &mut crate::providers::Usage, next: &crate::providers::Usage) {
|
||
total.prompt_tokens = total.prompt_tokens.saturating_add(next.prompt_tokens);
|
||
total.completion_tokens = total
|
||
.completion_tokens
|
||
.saturating_add(next.completion_tokens);
|
||
total.total_tokens = total.total_tokens.saturating_add(next.total_tokens);
|
||
total.cached_tokens = sum_optional_tokens(total.cached_tokens, next.cached_tokens);
|
||
total.cache_read_input_tokens =
|
||
sum_optional_tokens(total.cache_read_input_tokens, next.cache_read_input_tokens);
|
||
total.cache_creation_input_tokens = sum_optional_tokens(
|
||
total.cache_creation_input_tokens,
|
||
next.cache_creation_input_tokens,
|
||
);
|
||
}
|
||
|
||
fn sum_optional_tokens(left: Option<u32>, right: Option<u32>) -> Option<u32> {
|
||
match (left, right) {
|
||
(None, None) => None,
|
||
(left, right) => Some(
|
||
left.unwrap_or_default()
|
||
.saturating_add(right.unwrap_or_default()),
|
||
),
|
||
}
|
||
}
|
||
|
||
impl AgentLoop {
|
||
/// Create a new AgentLoop with a provider created from config.
|
||
pub fn new(provider_config: LLMProviderConfig) -> Result<Self, AgentError> {
|
||
let max_iterations = provider_config.max_tool_iterations;
|
||
let model_name = provider_config.model_id.clone();
|
||
let workspace_dir = provider_config.workspace_dir.clone();
|
||
let input_types = provider_config.input_types.clone();
|
||
let provider = create_provider(provider_config)
|
||
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
||
|
||
Ok(Self {
|
||
provider: Arc::from(provider),
|
||
tools: Arc::new(ToolRegistry::new()),
|
||
observer: None,
|
||
context_window: 0,
|
||
max_iterations,
|
||
workspace_dir,
|
||
model_name,
|
||
input_types,
|
||
media_registry: MediaHandlerRegistry::with_defaults(),
|
||
})
|
||
}
|
||
|
||
/// Create a new AgentLoop with provider created from config and given tools.
|
||
pub fn with_tools(
|
||
provider_config: LLMProviderConfig,
|
||
tools: Arc<ToolRegistry>,
|
||
) -> Result<Self, AgentError> {
|
||
let max_iterations = provider_config.max_tool_iterations;
|
||
let model_name = provider_config.model_id.clone();
|
||
let workspace_dir = provider_config.workspace_dir.clone();
|
||
let input_types = provider_config.input_types.clone();
|
||
let provider = create_provider(provider_config)
|
||
.map_err(|e| AgentError::ProviderCreation(e.to_string()))?;
|
||
|
||
Ok(Self {
|
||
provider: Arc::from(provider),
|
||
tools,
|
||
observer: None,
|
||
context_window: 0,
|
||
max_iterations,
|
||
workspace_dir,
|
||
model_name,
|
||
input_types,
|
||
media_registry: MediaHandlerRegistry::with_defaults(),
|
||
})
|
||
}
|
||
|
||
/// Create a new AgentLoop with an existing shared provider.
|
||
pub fn with_provider(
|
||
provider: Arc<dyn LLMProvider>,
|
||
max_iterations: usize,
|
||
model_name: String,
|
||
workspace_dir: PathBuf,
|
||
input_types: Vec<String>,
|
||
) -> Self {
|
||
Self {
|
||
provider,
|
||
tools: Arc::new(ToolRegistry::new()),
|
||
observer: None,
|
||
context_window: 0,
|
||
max_iterations,
|
||
workspace_dir,
|
||
model_name,
|
||
input_types,
|
||
media_registry: MediaHandlerRegistry::with_defaults(),
|
||
}
|
||
}
|
||
|
||
/// Create a new AgentLoop with an existing shared provider and given tools.
|
||
pub fn with_provider_and_tools(
|
||
provider: Arc<dyn LLMProvider>,
|
||
tools: Arc<ToolRegistry>,
|
||
max_iterations: usize,
|
||
model_name: String,
|
||
workspace_dir: PathBuf,
|
||
input_types: Vec<String>,
|
||
) -> Self {
|
||
Self {
|
||
provider,
|
||
tools,
|
||
observer: None,
|
||
context_window: 0,
|
||
max_iterations,
|
||
workspace_dir,
|
||
model_name,
|
||
input_types,
|
||
media_registry: MediaHandlerRegistry::with_defaults(),
|
||
}
|
||
}
|
||
|
||
/// Set the context window size for preemptive trimming.
|
||
pub fn with_context_window(mut self, window: usize) -> Self {
|
||
self.context_window = window;
|
||
self
|
||
}
|
||
|
||
/// Set the workspace directory.
|
||
pub fn with_workspace_dir(mut self, dir: PathBuf) -> Self {
|
||
self.workspace_dir = dir;
|
||
self
|
||
}
|
||
|
||
/// Set an observer for tracking events.
|
||
pub fn with_observer(mut self, observer: Arc<dyn Observer>) -> Self {
|
||
self.observer = Some(observer);
|
||
self
|
||
}
|
||
|
||
/// Preemptive trim: truncate old tool results in-place when history is
|
||
/// approaching the context window limit. Old results (outside of `keep_recent`
|
||
/// zone) are replaced with a short placeholder; recent results are truncated
|
||
/// to `max_chars`.
|
||
fn preemptive_trim_old_tool_results(
|
||
&self,
|
||
messages: &mut [ChatMessage],
|
||
max_chars: usize,
|
||
keep_recent: usize,
|
||
) -> usize {
|
||
let end = messages.len().saturating_sub(keep_recent);
|
||
let start = 1; // protect system message at [0] if present
|
||
let mut modified = 0;
|
||
for message in messages.iter_mut().take(end).skip(start) {
|
||
if message.role != "tool" {
|
||
continue;
|
||
}
|
||
if message.content.len() <= max_chars {
|
||
continue;
|
||
}
|
||
let tool_name = message.tool_name.as_deref().unwrap_or("unknown");
|
||
let chars = message.content.len();
|
||
message.content = format!(
|
||
"[Tool output ({}) — {} chars, omitted from context]",
|
||
tool_name, chars
|
||
);
|
||
modified += 1;
|
||
}
|
||
modified
|
||
}
|
||
|
||
pub fn tools(&self) -> &Arc<ToolRegistry> {
|
||
&self.tools
|
||
}
|
||
|
||
async fn stream_completion(
|
||
&self,
|
||
request: ChatCompletionRequest,
|
||
iteration: u32,
|
||
turn: Option<&AgentTurnContext>,
|
||
) -> Result<ChatCompletionResponse, AgentError> {
|
||
let metrics = crate::observability::metrics::global_metrics();
|
||
let provider_name = self.provider.name().to_string();
|
||
let provider_model = self.provider.model_id().to_string();
|
||
let start = Instant::now();
|
||
|
||
let mut provider_stream = match self.provider.stream(request).await {
|
||
Ok(stream) => stream,
|
||
Err(error) => {
|
||
tracing::error!(error = %error, "LLM request failed");
|
||
metrics.record_provider(
|
||
&provider_name,
|
||
&provider_model,
|
||
None,
|
||
start.elapsed().as_millis() as u64,
|
||
true,
|
||
);
|
||
return Err(AgentError::LlmError(error.to_string()));
|
||
}
|
||
};
|
||
let mut accumulator = ProviderResponseAccumulator::default();
|
||
while let Some(chunk) = provider_stream.next().await {
|
||
let chunk = match chunk {
|
||
Ok(chunk) => chunk,
|
||
Err(error) => {
|
||
tracing::error!(error = %error, "LLM stream failed");
|
||
metrics.record_provider(
|
||
&provider_name,
|
||
&provider_model,
|
||
None,
|
||
start.elapsed().as_millis() as u64,
|
||
true,
|
||
);
|
||
return Err(AgentError::LlmError(error.to_string()));
|
||
}
|
||
};
|
||
if let Some(turn) = turn {
|
||
let event = match &chunk {
|
||
ProviderChunk::Reasoning(delta) => Some(TurnEvent::ReasoningDelta {
|
||
iteration,
|
||
delta: delta.clone(),
|
||
}),
|
||
ProviderChunk::Text(delta) => Some(TurnEvent::TextDelta {
|
||
iteration,
|
||
delta: delta.clone(),
|
||
}),
|
||
_ => None,
|
||
};
|
||
if let Some(event) = event {
|
||
turn.emitter.emit(event).map_err(|error| {
|
||
AgentError::Other(format!("turn event rejected: {error}"))
|
||
})?;
|
||
}
|
||
}
|
||
accumulator.push(chunk);
|
||
}
|
||
let response = accumulator.finish();
|
||
let latency_ms = start.elapsed().as_millis() as u64;
|
||
metrics.record_provider(&provider_name, &provider_model, None, latency_ms, false);
|
||
metrics.record_provider_tokens(&provider_name, &response.usage);
|
||
Ok(response)
|
||
}
|
||
|
||
fn annotate_message(
|
||
message: &mut ChatMessage,
|
||
turn: Option<&AgentTurnContext>,
|
||
iteration: u32,
|
||
final_response: bool,
|
||
) {
|
||
message.iteration = Some(iteration);
|
||
if let Some(turn) = turn {
|
||
message.turn_id = Some(turn.turn_id.clone());
|
||
if final_response {
|
||
message.id = turn.message_id.clone();
|
||
}
|
||
}
|
||
}
|
||
|
||
fn chat_message_to_llm_message(&self, m: &ChatMessage, include_media: bool) -> Message {
|
||
let content = if m.media_refs.is_empty() || !include_media {
|
||
vec![ContentBlock::text(&m.content)]
|
||
} else {
|
||
let origin = match m.role.as_str() {
|
||
"user" => MediaOrigin::User,
|
||
"tool" => MediaOrigin::Tool(m.tool_name.as_deref().unwrap_or("unknown")),
|
||
_ => MediaOrigin::Message,
|
||
};
|
||
// Provider APIs generally allow native image/audio blocks only in
|
||
// user input (and, through provider-specific adaptation, current
|
||
// tool results). Persisted assistant attachments are delivery
|
||
// artifacts: replay their manifest as text, never as native media.
|
||
let native_input_types = if matches!(m.role.as_str(), "user" | "tool") {
|
||
self.input_types.as_slice()
|
||
} else {
|
||
&[]
|
||
};
|
||
build_content_blocks(
|
||
&m.content,
|
||
&m.media_refs,
|
||
native_input_types,
|
||
&self.media_registry,
|
||
origin,
|
||
)
|
||
};
|
||
|
||
Message {
|
||
role: m.role.clone(),
|
||
content,
|
||
reasoning_content: m.reasoning_content.clone(),
|
||
provider_state: m.provider_state.clone(),
|
||
tool_call_id: m.tool_call_id.clone(),
|
||
name: m.tool_name.clone(),
|
||
tool_calls: m.tool_calls.clone(),
|
||
}
|
||
}
|
||
|
||
fn messages_for_llm(&self, messages: &[ChatMessage]) -> Vec<Message> {
|
||
messages
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(index, message)| {
|
||
let include_media = should_include_message_media(messages, index);
|
||
self.chat_message_to_llm_message(message, include_media)
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Process a message using the provided conversation history.
|
||
/// History management is handled externally by SessionManager.
|
||
///
|
||
/// This method supports multi-round tool calling: after executing tools,
|
||
/// it loops back to the LLM with the tool results until either:
|
||
/// - The LLM returns no more tool calls (final response)
|
||
/// - Maximum iterations are reached
|
||
pub async fn process(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
) -> Result<AgentProcessResult, AgentError> {
|
||
self.process_inner(messages, None, ToolExecutionContext::default())
|
||
.await
|
||
}
|
||
|
||
pub async fn process_with_context(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
tool_context: ToolExecutionContext,
|
||
) -> Result<AgentProcessResult, AgentError> {
|
||
self.process_inner(messages, None, tool_context).await
|
||
}
|
||
|
||
pub async fn process_streaming(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
turn: AgentTurnContext,
|
||
) -> Result<AgentProcessResult, AgentError> {
|
||
let tool_context = ToolExecutionContext::default().with_turn_id(turn.turn_id.clone());
|
||
self.process_inner(messages, Some(turn), tool_context).await
|
||
}
|
||
|
||
pub async fn process_streaming_with_context(
|
||
&self,
|
||
messages: Vec<ChatMessage>,
|
||
turn: AgentTurnContext,
|
||
mut tool_context: ToolExecutionContext,
|
||
) -> Result<AgentProcessResult, AgentError> {
|
||
if tool_context.turn_id.is_none() {
|
||
tool_context.turn_id = Some(turn.turn_id.clone());
|
||
}
|
||
self.process_inner(messages, Some(turn), tool_context).await
|
||
}
|
||
|
||
async fn process_inner(
|
||
&self,
|
||
mut messages: Vec<ChatMessage>,
|
||
turn: Option<AgentTurnContext>,
|
||
tool_context: ToolExecutionContext,
|
||
) -> Result<AgentProcessResult, AgentError> {
|
||
let turn_start = Instant::now();
|
||
|
||
#[cfg(debug_assertions)]
|
||
tracing::debug!(
|
||
history_len = messages.len(),
|
||
max_iterations = self.max_iterations,
|
||
"Starting agent process"
|
||
);
|
||
|
||
// Build and inject system prompt if not present
|
||
let has_system = messages.first().is_some_and(|m| m.role == "system");
|
||
if !has_system {
|
||
let system_prompt =
|
||
build_system_prompt(&self.workspace_dir, &self.model_name, &self.tools);
|
||
#[cfg(debug_assertions)]
|
||
tracing::debug!("System prompt injected:\n{}", system_prompt);
|
||
messages.insert(0, ChatMessage::system(system_prompt));
|
||
}
|
||
|
||
// Track tool calls for loop detection
|
||
let mut loop_detector = LoopDetector::new(LoopDetectorConfig::default());
|
||
let mut emitted_messages = Vec::new();
|
||
let mut accumulated_tokens: u32 = 0;
|
||
let mut accumulated_usage = crate::providers::Usage::default();
|
||
let mut last_request_usage = None;
|
||
|
||
for iteration in 0..self.max_iterations {
|
||
#[cfg(debug_assertions)]
|
||
tracing::debug!(iteration, "Agent iteration started");
|
||
|
||
// Preemptive context check: trim old tool results if token estimate
|
||
// exceeds 80% of context window to prevent mid-loop overflow.
|
||
if self.context_window > 0 {
|
||
let estimated = estimate_tokens(&messages);
|
||
let danger = (self.context_window as f64 * 0.8) as usize;
|
||
if estimated > danger {
|
||
let trimmed = self.preemptive_trim_old_tool_results(&mut messages, 2000, 4);
|
||
if trimmed > 0 {
|
||
#[cfg(debug_assertions)]
|
||
tracing::debug!(
|
||
estimated,
|
||
danger,
|
||
trimmed_msgs = trimmed,
|
||
"Preemptive tool-result trim applied in loop"
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Convert messages to LLM format
|
||
let messages_for_llm = self.messages_for_llm(&messages);
|
||
|
||
// Build request
|
||
let tools = if self.tools.has_tools() {
|
||
Some(self.tools.get_definitions())
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let request = ChatCompletionRequest {
|
||
messages: messages_for_llm,
|
||
temperature: None,
|
||
max_tokens: None,
|
||
tools,
|
||
};
|
||
|
||
// Call LLM
|
||
let iteration = u32::try_from(iteration)
|
||
.map_err(|_| AgentError::Other("tool iteration exceeds u32".to_string()))?;
|
||
let response = self
|
||
.stream_completion(request, iteration, turn.as_ref())
|
||
.await?;
|
||
|
||
accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens);
|
||
merge_usage(&mut accumulated_usage, &response.usage);
|
||
last_request_usage = Some(response.usage.clone());
|
||
|
||
#[cfg(debug_assertions)]
|
||
tracing::debug!(
|
||
iteration,
|
||
response_len = response.content.len(),
|
||
tool_calls_len = response.tool_calls.len(),
|
||
"LLM response received"
|
||
);
|
||
|
||
// If no tool calls, this is the final response
|
||
if response.tool_calls.is_empty() {
|
||
let mut assistant_message = ChatMessage::assistant(response.content);
|
||
assistant_message.reasoning_content = response.reasoning_content;
|
||
assistant_message.provider_state = response.provider_state;
|
||
Self::annotate_message(&mut assistant_message, turn.as_ref(), iteration, true);
|
||
emitted_messages.push(assistant_message.clone());
|
||
crate::observability::metrics::global_metrics().record_turn(
|
||
Some(&accumulated_usage),
|
||
turn_start.elapsed().as_millis() as u64,
|
||
);
|
||
return Ok(AgentProcessResult {
|
||
final_response: assistant_message,
|
||
emitted_messages,
|
||
total_tokens: Some(accumulated_tokens),
|
||
usage: Some(accumulated_usage),
|
||
last_request_usage,
|
||
});
|
||
}
|
||
|
||
if let Some(turn) = turn.as_ref() {
|
||
turn.emitter
|
||
.emit(TurnEvent::TextSegmentFinished { iteration })
|
||
.map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?;
|
||
}
|
||
|
||
// Execute tool calls. User-visible progress is emitted through the
|
||
// structured TurnEvent stream, not a second notification channel.
|
||
{
|
||
let tools_info: Vec<String> = response
|
||
.tool_calls
|
||
.iter()
|
||
.map(|tc| {
|
||
let args = serde_json::to_string(&tc.arguments).unwrap_or_default();
|
||
let s = format!("{}:{}", tc.name, args);
|
||
s
|
||
})
|
||
.collect();
|
||
tracing::info!(iteration, count = response.tool_calls.len(), tools = %tools_info.join(", "), "Tool calls detected, executing tools");
|
||
}
|
||
|
||
// Add assistant message with tool calls
|
||
let mut assistant_message = ChatMessage::assistant_with_tool_calls(
|
||
response.content.clone(),
|
||
response.tool_calls.clone(),
|
||
);
|
||
assistant_message.reasoning_content = response.reasoning_content;
|
||
assistant_message.provider_state = response.provider_state;
|
||
Self::annotate_message(&mut assistant_message, turn.as_ref(), iteration, false);
|
||
messages.push(assistant_message.clone());
|
||
emitted_messages.push(assistant_message);
|
||
|
||
// Execute tools and add results to messages
|
||
let tool_results = self
|
||
.execute_tools(
|
||
&response.tool_calls,
|
||
iteration,
|
||
turn.as_ref(),
|
||
&tool_context,
|
||
)
|
||
.await?;
|
||
|
||
for (tool_call, result) in response.tool_calls.iter().zip(tool_results.iter()) {
|
||
// Log function call with name and arguments
|
||
let args_str = match &tool_call.arguments {
|
||
serde_json::Value::Object(obj) if obj.is_empty() => "{}".to_string(),
|
||
other => {
|
||
serde_json::to_string_pretty(other).unwrap_or_else(|_| other.to_string())
|
||
}
|
||
};
|
||
tracing::info!(tool = %tool_call.name, args = %args_str, "Calling tool");
|
||
|
||
// Truncate tool result if too large
|
||
let truncated_output = truncate_tool_result(&result.output);
|
||
|
||
// Record tool call and check for loops
|
||
let loop_result = loop_detector.record(&tool_call.name, &tool_call.arguments);
|
||
|
||
match loop_result {
|
||
LoopDetectionResult::Warning(msg) => {
|
||
// Add warning and proceed
|
||
tracing::warn!(
|
||
tool = %tool_call.name,
|
||
"Loop warning: {}",
|
||
msg
|
||
);
|
||
let mut tool_message = ChatMessage::tool_with_media(
|
||
tool_call.id.clone(),
|
||
tool_call.name.clone(),
|
||
format!("{}\n\n[上一条结果]\n{}", msg, truncated_output),
|
||
result.media_refs.clone(),
|
||
);
|
||
Self::annotate_message(&mut tool_message, turn.as_ref(), iteration, false);
|
||
messages.push(tool_message.clone());
|
||
emitted_messages.push(tool_message);
|
||
}
|
||
LoopDetectionResult::Ok => {
|
||
let mut tool_message = ChatMessage::tool_with_media(
|
||
tool_call.id.clone(),
|
||
tool_call.name.clone(),
|
||
truncated_output,
|
||
result.media_refs.clone(),
|
||
);
|
||
Self::annotate_message(&mut tool_message, turn.as_ref(), iteration, false);
|
||
messages.push(tool_message.clone());
|
||
emitted_messages.push(tool_message);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Loop continues to next iteration with updated messages
|
||
#[cfg(debug_assertions)]
|
||
tracing::debug!(
|
||
iteration,
|
||
message_count = messages.len(),
|
||
"Tool execution complete, continuing to next iteration"
|
||
);
|
||
}
|
||
|
||
// Max iterations reached - ask LLM for a summary based on completed work
|
||
tracing::warn!("Max iterations reached, requesting final summary from LLM");
|
||
|
||
// Add a message asking for summary
|
||
let summary_request = ChatMessage::user(
|
||
"You have reached the maximum number of tool call iterations. \
|
||
Please provide your best answer based on the work completed so far.",
|
||
);
|
||
messages.push(summary_request);
|
||
|
||
// Convert messages to LLM format
|
||
let messages_for_llm = self.messages_for_llm(&messages);
|
||
|
||
let request = ChatCompletionRequest {
|
||
messages: messages_for_llm,
|
||
temperature: None,
|
||
max_tokens: None,
|
||
tools: None, // No tools in final summary call
|
||
};
|
||
|
||
let summary_iteration = u32::try_from(self.max_iterations)
|
||
.map_err(|_| AgentError::Other("tool iteration exceeds u32".to_string()))?;
|
||
match self
|
||
.stream_completion(request, summary_iteration, turn.as_ref())
|
||
.await
|
||
{
|
||
Ok(response) => {
|
||
accumulated_tokens = accumulated_tokens.saturating_add(response.usage.total_tokens);
|
||
merge_usage(&mut accumulated_usage, &response.usage);
|
||
last_request_usage = Some(response.usage.clone());
|
||
let mut assistant_message = ChatMessage::assistant(response.content);
|
||
assistant_message.reasoning_content = response.reasoning_content;
|
||
assistant_message.provider_state = response.provider_state;
|
||
Self::annotate_message(
|
||
&mut assistant_message,
|
||
turn.as_ref(),
|
||
summary_iteration,
|
||
true,
|
||
);
|
||
emitted_messages.push(assistant_message.clone());
|
||
crate::observability::metrics::global_metrics().record_turn(
|
||
Some(&accumulated_usage),
|
||
turn_start.elapsed().as_millis() as u64,
|
||
);
|
||
Ok(AgentProcessResult {
|
||
final_response: assistant_message,
|
||
emitted_messages,
|
||
total_tokens: Some(accumulated_tokens),
|
||
usage: Some(accumulated_usage),
|
||
last_request_usage,
|
||
})
|
||
}
|
||
Err(e) => {
|
||
// Fallback if summary call fails
|
||
tracing::error!(error = %e, "Failed to get summary from LLM");
|
||
let fallback = format!(
|
||
"I reached the maximum number of tool call iterations ({}) without completing the task. The work done so far has been lost due to an error. Please try breaking the task into smaller steps.",
|
||
self.max_iterations
|
||
);
|
||
if let Some(turn) = turn.as_ref() {
|
||
turn.emitter
|
||
.emit(TurnEvent::TextSegmentFinished {
|
||
iteration: summary_iteration,
|
||
})
|
||
.and_then(|()| {
|
||
turn.emitter.emit(TurnEvent::TextDelta {
|
||
iteration: summary_iteration,
|
||
delta: fallback.clone(),
|
||
})
|
||
})
|
||
.map_err(|error| {
|
||
AgentError::Other(format!("turn event rejected: {error}"))
|
||
})?;
|
||
}
|
||
let mut final_message = ChatMessage::assistant(fallback);
|
||
Self::annotate_message(&mut final_message, turn.as_ref(), summary_iteration, true);
|
||
emitted_messages.push(final_message.clone());
|
||
let turn_usage = (accumulated_usage.total_tokens > 0).then_some(&accumulated_usage);
|
||
crate::observability::metrics::global_metrics()
|
||
.record_turn(turn_usage, turn_start.elapsed().as_millis() as u64);
|
||
Ok(AgentProcessResult {
|
||
final_response: final_message,
|
||
emitted_messages,
|
||
total_tokens: if accumulated_tokens > 0 {
|
||
Some(accumulated_tokens)
|
||
} else {
|
||
None
|
||
},
|
||
usage: (accumulated_usage.total_tokens > 0).then_some(accumulated_usage),
|
||
last_request_usage,
|
||
})
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Determine whether to execute tools in parallel or sequentially.
|
||
///
|
||
/// Returns true if:
|
||
/// - There are multiple tool calls
|
||
/// - None of the tools require sequential execution (tool_search, non-concurrency-safe)
|
||
fn should_execute_in_parallel(&self, tool_calls: &[ToolCall]) -> bool {
|
||
if tool_calls.len() <= 1 {
|
||
return false;
|
||
}
|
||
|
||
// tool_search must run sequentially to avoid MCP activation race conditions
|
||
if tool_calls.iter().any(|tc| tc.name == "tool_search") {
|
||
return false;
|
||
}
|
||
|
||
// All tools must be concurrency-safe to run in parallel
|
||
tool_calls.iter().all(|tc| {
|
||
self.tools
|
||
.get(&tc.name)
|
||
.map(|t| t.concurrency_safe())
|
||
.unwrap_or(false)
|
||
})
|
||
}
|
||
|
||
/// Execute multiple tool calls, choosing parallel or sequential based on conditions.
|
||
async fn execute_tools(
|
||
&self,
|
||
tool_calls: &[ToolCall],
|
||
iteration: u32,
|
||
turn: Option<&AgentTurnContext>,
|
||
context: &ToolExecutionContext,
|
||
) -> Result<Vec<ToolExecutionOutcome>, AgentError> {
|
||
if self.should_execute_in_parallel(tool_calls) {
|
||
tracing::debug!("Executing {} tools in parallel", tool_calls.len());
|
||
self.execute_tools_parallel(tool_calls, iteration, turn, context)
|
||
.await
|
||
} else {
|
||
tracing::debug!("Executing {} tools sequentially", tool_calls.len());
|
||
self.execute_tools_sequential(tool_calls, iteration, turn, context)
|
||
.await
|
||
}
|
||
}
|
||
|
||
/// Execute tools in parallel using join_all.
|
||
async fn execute_tools_parallel(
|
||
&self,
|
||
tool_calls: &[ToolCall],
|
||
iteration: u32,
|
||
turn: Option<&AgentTurnContext>,
|
||
context: &ToolExecutionContext,
|
||
) -> Result<Vec<ToolExecutionOutcome>, AgentError> {
|
||
let futures: Vec<_> = tool_calls
|
||
.iter()
|
||
.map(|tool_call| self.execute_one_tool(tool_call, iteration, turn, context))
|
||
.collect();
|
||
|
||
futures_util::future::join_all(futures)
|
||
.await
|
||
.into_iter()
|
||
.collect()
|
||
}
|
||
|
||
/// Execute tools sequentially.
|
||
async fn execute_tools_sequential(
|
||
&self,
|
||
tool_calls: &[ToolCall],
|
||
iteration: u32,
|
||
turn: Option<&AgentTurnContext>,
|
||
context: &ToolExecutionContext,
|
||
) -> Result<Vec<ToolExecutionOutcome>, AgentError> {
|
||
let mut outcomes = Vec::with_capacity(tool_calls.len());
|
||
|
||
for tool_call in tool_calls {
|
||
outcomes.push(
|
||
self.execute_one_tool(tool_call, iteration, turn, context)
|
||
.await?,
|
||
);
|
||
}
|
||
|
||
Ok(outcomes)
|
||
}
|
||
|
||
/// Execute a single tool and return the outcome with event tracking.
|
||
async fn execute_one_tool(
|
||
&self,
|
||
tool_call: &ToolCall,
|
||
iteration: u32,
|
||
turn: Option<&AgentTurnContext>,
|
||
context: &ToolExecutionContext,
|
||
) -> Result<ToolExecutionOutcome, AgentError> {
|
||
let start = Instant::now();
|
||
let tool_name = tool_call.name.clone();
|
||
|
||
if let Some(turn) = turn {
|
||
turn.emitter
|
||
.emit(TurnEvent::ToolStarted {
|
||
iteration,
|
||
call: tool_call.clone(),
|
||
})
|
||
.map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?;
|
||
}
|
||
|
||
// Record ToolCallStart event
|
||
if let Some(ref observer) = self.observer {
|
||
observer.record_event(&ObserverEvent::ToolCallStart {
|
||
tool: tool_name.clone(),
|
||
arguments: Some(truncate_args(&tool_call.arguments, 300)),
|
||
});
|
||
}
|
||
|
||
let result = self.execute_tool_internal(tool_call, context).await;
|
||
let duration = start.elapsed();
|
||
|
||
if let Some(turn) = turn {
|
||
turn.emitter
|
||
.emit(TurnEvent::ToolFinished {
|
||
iteration,
|
||
call_id: tool_call.id.clone(),
|
||
success: result.success,
|
||
preview: Some(tool_result_preview(&truncate_tool_result(&result.output))),
|
||
})
|
||
.map_err(|error| AgentError::Other(format!("turn event rejected: {error}")))?;
|
||
}
|
||
|
||
// Record ToolCall event
|
||
if let Some(ref observer) = self.observer {
|
||
observer.record_event(&ObserverEvent::ToolCall {
|
||
tool: tool_name.clone(),
|
||
duration,
|
||
success: result.success,
|
||
});
|
||
}
|
||
|
||
crate::observability::metrics::global_metrics()
|
||
.record_tool_call(&tool_name, result.success);
|
||
|
||
// Apply duration
|
||
Ok(ToolExecutionOutcome { duration, ..result })
|
||
}
|
||
|
||
/// Internal tool execution without event tracking.
|
||
async fn execute_tool_internal(
|
||
&self,
|
||
tool_call: &ToolCall,
|
||
context: &ToolExecutionContext,
|
||
) -> ToolExecutionOutcome {
|
||
let tool = match self.tools.get(&tool_call.name) {
|
||
Some(t) => t,
|
||
None => {
|
||
tracing::warn!(tool = %tool_call.name, "Tool not found");
|
||
return ToolExecutionOutcome::failure(
|
||
format!("Error: Tool '{}' not found", tool_call.name),
|
||
Some(format!("Tool '{}' not found", tool_call.name)),
|
||
);
|
||
}
|
||
};
|
||
|
||
match tool
|
||
.execute_with_context(context, tool_call.arguments.clone())
|
||
.await
|
||
{
|
||
Ok(result_with_media) => {
|
||
let result = result_with_media.result;
|
||
if result.success {
|
||
ToolExecutionOutcome::success_with_media(
|
||
result.output,
|
||
result_with_media.media_refs,
|
||
)
|
||
} else {
|
||
let error = result.error.unwrap_or_default();
|
||
ToolExecutionOutcome::failure(format!("Error: {}", error), Some(error))
|
||
}
|
||
}
|
||
Err(e) => {
|
||
tracing::error!(tool = %tool_call.name, error = %e, "Tool execution failed");
|
||
ToolExecutionOutcome::failure(format!("Error: {}", e), Some(e.to_string()))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::observability::{MultiObserver, Observer};
|
||
use crate::providers::{
|
||
ChatCompletionResponse, FinishReason, ProviderChunk, ProviderStream, Usage,
|
||
};
|
||
use crate::session::{TurnBlock, TurnController};
|
||
use crate::tools::FileReadTool;
|
||
|
||
struct TestObserver {
|
||
events: std::sync::Mutex<Vec<ObserverEvent>>,
|
||
}
|
||
|
||
struct StreamingTextProvider;
|
||
|
||
#[async_trait::async_trait]
|
||
impl LLMProvider for StreamingTextProvider {
|
||
async fn stream(
|
||
&self,
|
||
_request: ChatCompletionRequest,
|
||
) -> Result<ProviderStream, crate::providers::DynProviderError> {
|
||
let chunks = vec![
|
||
ProviderChunk::Metadata {
|
||
id: "response".into(),
|
||
model: "streaming-test".into(),
|
||
},
|
||
ProviderChunk::Reasoning("because ".into()),
|
||
ProviderChunk::Reasoning("facts".into()),
|
||
ProviderChunk::Text("hello ".into()),
|
||
ProviderChunk::Text("world".into()),
|
||
ProviderChunk::ProviderState(crate::bus::ProviderReasoningState {
|
||
provider: "test".into(),
|
||
payload: serde_json::json!({"opaque":"state"}),
|
||
}),
|
||
ProviderChunk::Usage(Usage {
|
||
prompt_tokens: 2,
|
||
completion_tokens: 3,
|
||
total_tokens: 5,
|
||
..Usage::default()
|
||
}),
|
||
ProviderChunk::Done(FinishReason::Stop),
|
||
];
|
||
Ok(Box::pin(futures_util::stream::iter(
|
||
chunks.into_iter().map(Ok),
|
||
)))
|
||
}
|
||
|
||
fn ptype(&self) -> &str {
|
||
"test"
|
||
}
|
||
|
||
fn name(&self) -> &str {
|
||
"streaming-test"
|
||
}
|
||
|
||
fn model_id(&self) -> &str {
|
||
"streaming-test"
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn process_streaming_emits_turn_blocks_and_stamps_durable_message() {
|
||
let agent = AgentLoop::with_provider(
|
||
Arc::new(StreamingTextProvider),
|
||
1,
|
||
"streaming-test".into(),
|
||
PathBuf::from("."),
|
||
Vec::new(),
|
||
);
|
||
let (controller, emitter, _) = TurnController::start("session", "assistant-id");
|
||
let initial = controller.snapshot();
|
||
let context =
|
||
AgentTurnContext::new(initial.id.0.clone(), initial.message_id.clone(), emitter);
|
||
|
||
let result = agent
|
||
.process_streaming(vec![ChatMessage::user("hi")], context)
|
||
.await
|
||
.unwrap();
|
||
|
||
assert_eq!(result.final_response.id, "assistant-id");
|
||
assert_eq!(result.final_response.content, "hello world");
|
||
assert_eq!(
|
||
result.final_response.reasoning_content.as_deref(),
|
||
Some("because facts")
|
||
);
|
||
assert_eq!(result.final_response.turn_id, Some(initial.id.0.clone()));
|
||
assert_eq!(result.final_response.iteration, Some(0));
|
||
assert_eq!(
|
||
result
|
||
.final_response
|
||
.provider_state
|
||
.as_ref()
|
||
.map(|state| state.provider.as_str()),
|
||
Some("test")
|
||
);
|
||
assert_eq!(
|
||
result.usage.as_ref().map(|usage| usage.total_tokens),
|
||
Some(5)
|
||
);
|
||
|
||
let snapshot = controller.snapshot();
|
||
assert_eq!(snapshot.blocks.len(), 2);
|
||
assert!(matches!(
|
||
&snapshot.blocks[0],
|
||
TurnBlock::Reasoning { text, .. } if text == "because facts"
|
||
));
|
||
assert!(matches!(
|
||
&snapshot.blocks[1],
|
||
TurnBlock::Assistant { text, .. } if text == "hello world"
|
||
));
|
||
}
|
||
|
||
impl TestObserver {
|
||
fn new() -> Self {
|
||
Self {
|
||
events: std::sync::Mutex::new(Vec::new()),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl Observer for TestObserver {
|
||
fn record_event(&self, event: &ObserverEvent) {
|
||
self.events.lock().unwrap().push(event.clone());
|
||
}
|
||
|
||
fn name(&self) -> &str {
|
||
"test_observer"
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_observer_receives_tool_events() {
|
||
// Verify MultiObserver works
|
||
let mut multi = MultiObserver::new();
|
||
multi.add_observer(Box::new(TestObserver::new()));
|
||
|
||
let event = ObserverEvent::ToolCallStart {
|
||
tool: "test".to_string(),
|
||
arguments: Some("{}".to_string()),
|
||
};
|
||
multi.record_event(&event);
|
||
|
||
// Just verify the structure works
|
||
assert_eq!(multi.len(), 1);
|
||
}
|
||
|
||
struct ToolMediaProvider {
|
||
image_path: String,
|
||
requests: std::sync::Mutex<Vec<ChatCompletionRequest>>,
|
||
}
|
||
|
||
#[async_trait::async_trait]
|
||
impl LLMProvider for ToolMediaProvider {
|
||
async fn stream(
|
||
&self,
|
||
request: ChatCompletionRequest,
|
||
) -> Result<crate::providers::ProviderStream, crate::providers::DynProviderError> {
|
||
let call_number = {
|
||
let mut requests = self.requests.lock().unwrap();
|
||
requests.push(request);
|
||
requests.len()
|
||
};
|
||
let response = ChatCompletionResponse {
|
||
id: format!("response-{call_number}"),
|
||
model: "vision-test".to_string(),
|
||
content: if call_number == 1 {
|
||
String::new()
|
||
} else {
|
||
"image seen".to_string()
|
||
},
|
||
reasoning_content: None,
|
||
provider_state: None,
|
||
tool_calls: if call_number == 1 {
|
||
vec![ToolCall {
|
||
id: "call-image".to_string(),
|
||
name: "file_read".to_string(),
|
||
arguments: serde_json::json!({ "path": self.image_path }),
|
||
}]
|
||
} else {
|
||
Vec::new()
|
||
},
|
||
usage: Usage {
|
||
prompt_tokens: 1,
|
||
completion_tokens: 1,
|
||
total_tokens: 2,
|
||
cached_tokens: None,
|
||
cache_read_input_tokens: None,
|
||
cache_creation_input_tokens: None,
|
||
},
|
||
};
|
||
Ok(crate::providers::provider_stream_for_test(response))
|
||
}
|
||
|
||
fn ptype(&self) -> &str {
|
||
"test"
|
||
}
|
||
|
||
fn name(&self) -> &str {
|
||
"tool-media-test"
|
||
}
|
||
|
||
fn model_id(&self) -> &str {
|
||
"vision-test"
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn file_read_media_reaches_the_next_model_iteration() {
|
||
use std::io::Write;
|
||
|
||
let mut image = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
|
||
image.write_all(b"\x89PNG\r\n\x1a\nminimal").unwrap();
|
||
let provider = Arc::new(ToolMediaProvider {
|
||
image_path: image.path().to_string_lossy().into_owned(),
|
||
requests: std::sync::Mutex::new(Vec::new()),
|
||
});
|
||
let tools = Arc::new(ToolRegistry::new());
|
||
tools.register(FileReadTool::new());
|
||
let agent = AgentLoop::with_provider_and_tools(
|
||
provider.clone(),
|
||
tools,
|
||
2,
|
||
"vision-test".to_string(),
|
||
std::env::current_dir().unwrap(),
|
||
vec!["text".to_string(), "image".to_string()],
|
||
);
|
||
|
||
let (controller, emitter, _) = TurnController::start("session", "assistant-id");
|
||
let turn = controller.snapshot();
|
||
let result = agent
|
||
.process_streaming(
|
||
vec![ChatMessage::user("inspect the image")],
|
||
AgentTurnContext::new(turn.id.0.clone(), turn.message_id.clone(), emitter),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
assert_eq!(result.final_response.content, "image seen");
|
||
let requests = provider.requests.lock().unwrap();
|
||
assert_eq!(requests.len(), 2);
|
||
let tool_result = requests[1]
|
||
.messages
|
||
.iter()
|
||
.find(|message| message.role == "tool")
|
||
.unwrap();
|
||
assert!(
|
||
tool_result
|
||
.content
|
||
.iter()
|
||
.any(|block| matches!(block, ContentBlock::ImageUrl { .. }))
|
||
);
|
||
assert!(result.emitted_messages.iter().any(|message| {
|
||
message.role == "tool"
|
||
&& message
|
||
.media_refs
|
||
.iter()
|
||
.any(|media| media.media_type == "image")
|
||
}));
|
||
assert!(controller.snapshot().blocks.iter().any(|block| matches!(
|
||
block,
|
||
TurnBlock::Tool {
|
||
id,
|
||
status: crate::session::ToolStatus::Completed,
|
||
..
|
||
} if id == "call-image"
|
||
)));
|
||
}
|
||
|
||
#[test]
|
||
fn test_should_execute_in_parallel_single_tool() {
|
||
// Would need a proper setup with AgentLoop to test fully
|
||
// For now, just verify the logic: single tool should return false
|
||
let calls = [ToolCall {
|
||
id: "1".to_string(),
|
||
name: "test".to_string(),
|
||
arguments: serde_json::json!({}),
|
||
}];
|
||
|
||
// If there's only 1 tool, should return false regardless
|
||
assert!(calls.len() <= 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_chat_message_to_llm_message_preserves_assistant_tool_calls() {
|
||
use crate::providers::Message;
|
||
|
||
let chat_message = ChatMessage::assistant_with_tool_calls(
|
||
"calling tool",
|
||
vec![ToolCall {
|
||
id: "call_1".to_string(),
|
||
name: "calculator".to_string(),
|
||
arguments: serde_json::json!({ "expression": "2+2" }),
|
||
}],
|
||
);
|
||
|
||
let content = vec![ContentBlock::text(&chat_message.content)];
|
||
let provider_message = Message {
|
||
role: chat_message.role.clone(),
|
||
content,
|
||
reasoning_content: None,
|
||
provider_state: None,
|
||
tool_call_id: chat_message.tool_call_id.clone(),
|
||
name: chat_message.tool_name.clone(),
|
||
tool_calls: chat_message.tool_calls.clone(),
|
||
};
|
||
|
||
assert_eq!(provider_message.role, "assistant");
|
||
assert_eq!(provider_message.tool_calls.as_ref().unwrap().len(), 1);
|
||
assert_eq!(
|
||
provider_message.tool_calls.as_ref().unwrap()[0].id,
|
||
"call_1"
|
||
);
|
||
assert_eq!(
|
||
provider_message.tool_calls.as_ref().unwrap()[0].name,
|
||
"calculator"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_content_blocks_keeps_text_with_media() {
|
||
let registry = MediaHandlerRegistry::new();
|
||
let blocks = build_content_blocks(
|
||
"先看这段文字",
|
||
&[MediaRef {
|
||
path: "missing.png".to_string(),
|
||
media_type: "image".to_string(),
|
||
}],
|
||
&[],
|
||
®istry,
|
||
MediaOrigin::User,
|
||
);
|
||
|
||
assert_eq!(blocks.len(), 1);
|
||
assert!(matches!(blocks.first(), Some(ContentBlock::Text { text })
|
||
if text.starts_with("先看这段文字\n\n")
|
||
&& text.contains("随本条用户消息同时提交的附件")
|
||
&& text.contains("missing.png")
|
||
&& text.contains("\"media_type\":\"image\"")
|
||
&& text.contains("content is not embedded in this model request")));
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_content_blocks_describes_attachment_only_message() {
|
||
let registry = MediaHandlerRegistry::new();
|
||
let blocks = build_content_blocks(
|
||
"",
|
||
&[MediaRef {
|
||
path: "/tmp/report.docx".to_string(),
|
||
media_type: "file".to_string(),
|
||
}],
|
||
&[],
|
||
®istry,
|
||
MediaOrigin::User,
|
||
);
|
||
|
||
assert_eq!(blocks.len(), 1);
|
||
assert!(matches!(blocks.first(), Some(ContentBlock::Text { text })
|
||
if text.starts_with("[用户发送了以下附件")
|
||
&& text.contains("report.docx")
|
||
&& text.contains("application/vnd.openxmlformats-officedocument.wordprocessingml.document")
|
||
&& text.contains("\"extension\":\"docx\"")
|
||
&& text.contains("\"size_bytes\":null")
|
||
&& text.contains("content is not embedded in this model request")));
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_content_blocks_includes_path_for_supported_images() {
|
||
use std::io::Write;
|
||
|
||
let mut image = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
|
||
image.write_all(b"\x89PNG\r\n\x1a\nminimal").unwrap();
|
||
let path = image.path().to_string_lossy().into_owned();
|
||
let registry = MediaHandlerRegistry::with_defaults();
|
||
|
||
let blocks = build_content_blocks(
|
||
"分析图片",
|
||
&[MediaRef {
|
||
path: path.clone(),
|
||
media_type: "image".to_string(),
|
||
}],
|
||
&["image".to_string()],
|
||
®istry,
|
||
MediaOrigin::User,
|
||
);
|
||
|
||
assert!(matches!(blocks.first(), Some(ContentBlock::Text { text })
|
||
if text.starts_with("分析图片\n\n")
|
||
&& text.contains("随本条用户消息同时提交的附件")
|
||
&& text.contains(&path)
|
||
&& text.contains("model-native content block")));
|
||
assert!(matches!(blocks.get(1), Some(ContentBlock::ImageUrl { .. })));
|
||
}
|
||
|
||
#[test]
|
||
fn test_build_content_blocks_labels_tool_media_without_user_wording() {
|
||
use std::io::Write;
|
||
|
||
let mut image = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
|
||
image.write_all(b"\x89PNG\r\n\x1a\nminimal").unwrap();
|
||
let path = image.path().to_string_lossy().into_owned();
|
||
let registry = MediaHandlerRegistry::with_defaults();
|
||
|
||
let blocks = build_content_blocks(
|
||
"Image file ready for visual inspection.",
|
||
&[MediaRef {
|
||
path,
|
||
media_type: "image".to_string(),
|
||
}],
|
||
&["image".to_string()],
|
||
®istry,
|
||
MediaOrigin::Tool("file_read"),
|
||
);
|
||
|
||
assert!(matches!(blocks.first(), Some(ContentBlock::Text { text })
|
||
if text.contains("工具 file_read 返回了以下附件")
|
||
&& !text.contains("用户消息")));
|
||
assert!(matches!(blocks.get(1), Some(ContentBlock::ImageUrl { .. })));
|
||
}
|
||
|
||
#[test]
|
||
fn assistant_attachments_replay_as_text_without_native_image_blocks() {
|
||
use std::io::Write;
|
||
|
||
let mut image = tempfile::Builder::new().suffix(".png").tempfile().unwrap();
|
||
image.write_all(b"\x89PNG\r\n\x1a\nminimal").unwrap();
|
||
let path = image.path().to_string_lossy().into_owned();
|
||
let provider = Arc::new(ToolMediaProvider {
|
||
image_path: path.clone(),
|
||
requests: std::sync::Mutex::new(Vec::new()),
|
||
});
|
||
let agent = AgentLoop::with_provider_and_tools(
|
||
provider,
|
||
Arc::new(ToolRegistry::new()),
|
||
1,
|
||
"vision-test".to_string(),
|
||
std::env::current_dir().unwrap(),
|
||
vec!["text".to_string(), "image".to_string()],
|
||
);
|
||
let mut assistant = ChatMessage::assistant("截图已发送");
|
||
assistant.media_refs = vec![MediaRef {
|
||
path: path.clone(),
|
||
media_type: "image".to_string(),
|
||
}];
|
||
|
||
let converted = agent.chat_message_to_llm_message(&assistant, true);
|
||
|
||
assert_eq!(converted.role, "assistant");
|
||
assert_eq!(converted.content.len(), 1);
|
||
assert!(
|
||
matches!(converted.content.first(), Some(ContentBlock::Text { text })
|
||
if text.contains("截图已发送")
|
||
&& text.contains(&path)
|
||
&& text.contains("not embedded in this model request"))
|
||
);
|
||
assert!(
|
||
!converted
|
||
.content
|
||
.iter()
|
||
.any(|block| matches!(block, ContentBlock::ImageUrl { .. }))
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn only_the_trailing_tool_batch_replays_tool_media() {
|
||
let mut messages = vec![
|
||
ChatMessage::assistant_with_tool_calls(
|
||
"",
|
||
vec![ToolCall {
|
||
id: "old-call".to_string(),
|
||
name: "file_read".to_string(),
|
||
arguments: serde_json::json!({}),
|
||
}],
|
||
),
|
||
ChatMessage::tool_with_media(
|
||
"old-call",
|
||
"file_read",
|
||
"old",
|
||
vec![MediaRef {
|
||
path: "/tmp/old.png".to_string(),
|
||
media_type: "image".to_string(),
|
||
}],
|
||
),
|
||
ChatMessage::assistant("continue"),
|
||
ChatMessage::assistant_with_tool_calls(
|
||
"",
|
||
vec![ToolCall {
|
||
id: "new-call".to_string(),
|
||
name: "file_read".to_string(),
|
||
arguments: serde_json::json!({}),
|
||
}],
|
||
),
|
||
ChatMessage::tool_with_media(
|
||
"new-call",
|
||
"file_read",
|
||
"new",
|
||
vec![MediaRef {
|
||
path: "/tmp/new.png".to_string(),
|
||
media_type: "image".to_string(),
|
||
}],
|
||
),
|
||
];
|
||
|
||
assert!(!should_include_message_media(&messages, 1));
|
||
assert!(should_include_message_media(&messages, 4));
|
||
|
||
messages.push(ChatMessage::user("next turn"));
|
||
assert!(!should_include_message_media(&messages, 4));
|
||
}
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
pub enum AgentError {
|
||
ProviderCreation(String),
|
||
LlmError(String),
|
||
Other(String),
|
||
}
|
||
|
||
impl std::fmt::Display for AgentError {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
match self {
|
||
AgentError::ProviderCreation(e) => write!(f, "Provider creation error: {}", e),
|
||
AgentError::LlmError(e) => write!(f, "LLM error: {}", e),
|
||
AgentError::Other(e) => write!(f, "{}", e),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl std::error::Error for AgentError {}
|