use async_trait::async_trait;
use autoagents_core::agent::memory::SlidingWindowMemory;
use autoagents_core::agent::prebuilt::executor::ReActAgent;
use autoagents_core::agent::task::Task;
use autoagents_core::agent::{
ActorAgent, ActorAgentHandle, AgentBuilder, AgentDeriveT, AgentHooks, DirectAgent,
DirectAgentHandle,
};
use autoagents_core::runtime::{Runtime, SingleThreadedRuntime};
use autoagents_core::tool::{ToolT, to_llm_tool};
use autoagents_llm::chat::{
ChatMessage, ChatProvider, ChatResponse, ChatRole, MessageType, StructuredOutputFormat, Tool,
};
use autoagents_llm::completion::{CompletionProvider, CompletionRequest, CompletionResponse};
use autoagents_llm::embedding::EmbeddingProvider;
use autoagents_llm::error::LLMError;
use autoagents_llm::models::ModelsProvider;
use autoagents_llm::{FunctionCall, LLMProvider, ToolCall};
use autoagents_protocol::Event;
use criterion::{Criterion, criterion_group, criterion_main};
use futures_util::StreamExt;
use serde::Serialize;
use serde_json::{Value, json};
use std::hint::black_box;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Instant;
use tokio::runtime::Runtime as TokioRuntime;
use tokio::time::{Duration, sleep};
#[derive(Debug, Clone)]
struct BenchAgent;
const BENCH_TOOL_NAME: &str = "bench_tool";
const BENCH_TOOL_CALL_ID: &str = "bench_tool_call";
const DEFAULT_TOOL_ARGS: &str = "{}";
const DEFAULT_LLM_DELAY_MS: u64 = 5;
const DEFAULT_RESPONSE_TEXT: &str = "ok";
const DEFAULT_RUNTIME_BUFFER: usize = 1024;
#[derive(Debug, Clone)]
struct BenchTool;
#[async_trait]
impl autoagents_core::tool::ToolRuntime for BenchTool {
async fn execute(&self, args: Value) -> Result<Value, autoagents_core::tool::ToolCallError> {
Ok(json!({"ok": true, "args": args}))
}
}
impl ToolT for BenchTool {
fn name(&self) -> &str {
BENCH_TOOL_NAME
}
fn description(&self) -> &str {
"Benchmark tool that echoes inputs"
}
fn args_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"input": {"type": "string"}
},
"required": []
})
}
}
fn default_tool_call() -> ToolCall {
ToolCall {
id: BENCH_TOOL_CALL_ID.to_string(),
call_type: "function".to_string(),
function: FunctionCall {
name: BENCH_TOOL_NAME.to_string(),
arguments: DEFAULT_TOOL_ARGS.to_string(),
},
}
}
fn bench_tool_defs() -> Vec<Tool> {
let tool: Box<dyn ToolT> = Box::new(BenchTool);
vec![to_llm_tool(&tool)]
}
fn bench_tool_result_message() -> ChatMessage {
ChatMessage {
role: ChatRole::Tool,
message_type: MessageType::ToolResult(vec![default_tool_call()]),
content: String::default(),
}
}
async fn await_task_complete(
events: &mut autoagents_core::utils::BoxEventStream<Event>,
submission_id: autoagents_protocol::SubmissionId,
) -> Result<(), String> {
loop {
match events.next().await {
Some(Event::TaskComplete { sub_id, .. }) if sub_id == submission_id => return Ok(()),
Some(Event::TaskError { sub_id, error, .. }) if sub_id == submission_id => {
return Err(error);
}
Some(_) => {}
None => return Err("Event stream closed before task completion".to_string()),
}
}
}
#[derive(Debug, Clone)]
struct MockLlmProvider {
delay: Duration,
response_text: String,
}
impl MockLlmProvider {
fn new(delay_ms: u64) -> Self {
Self {
delay: Duration::from_millis(delay_ms),
response_text: DEFAULT_RESPONSE_TEXT.to_string(),
}
}
async fn simulate_latency(&self) {
if self.delay.as_nanos() > 0 {
sleep(self.delay).await;
}
}
fn should_return_tool_call(&self, messages: &[ChatMessage], tools: Option<&[Tool]>) -> bool {
if tools.is_none() {
return false;
}
!messages
.iter()
.any(|message| matches!(message.role, ChatRole::Tool))
}
fn build_tool_calls(&self, tools: Option<&[Tool]>) -> Option<Vec<ToolCall>> {
let tools = tools?;
let tool = tools.first()?;
let mut call = default_tool_call();
call.function.name = tool.function.name.clone();
Some(vec![call])
}
}
#[async_trait]
impl ChatProvider for MockLlmProvider {
async fn chat_with_tools(
&self,
messages: &[ChatMessage],
tools: Option<&[Tool]>,
_json_schema: Option<StructuredOutputFormat>,
) -> Result<Box<dyn ChatResponse>, LLMError> {
self.simulate_latency().await;
let tool_calls = if self.should_return_tool_call(messages, tools) {
self.build_tool_calls(tools)
} else {
None
};
Ok(Box::new(MockChatResponse {
text: Some(self.response_text.clone()),
tool_calls,
}))
}
}
#[async_trait]
impl CompletionProvider for MockLlmProvider {
async fn complete(
&self,
_req: &CompletionRequest,
_json_schema: Option<StructuredOutputFormat>,
) -> Result<CompletionResponse, LLMError> {
self.simulate_latency().await;
Ok(CompletionResponse {
text: self.response_text.clone(),
})
}
}
#[async_trait]
impl EmbeddingProvider for MockLlmProvider {
async fn embed(&self, input: Vec<String>) -> Result<Vec<Vec<f32>>, LLMError> {
self.simulate_latency().await;
Ok(input.iter().map(|_| vec![0.0, 1.0]).collect())
}
}
#[async_trait]
impl ModelsProvider for MockLlmProvider {}
impl LLMProvider for MockLlmProvider {}
#[derive(Debug)]
struct MockChatResponse {
text: Option<String>,
tool_calls: Option<Vec<ToolCall>>,
}
impl ChatResponse for MockChatResponse {
fn text(&self) -> Option<String> {
self.text.clone()
}
fn tool_calls(&self) -> Option<Vec<ToolCall>> {
self.tool_calls.clone()
}
}
impl std::fmt::Display for MockChatResponse {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.text.as_deref().unwrap_or(""))
}
}
impl AgentDeriveT for BenchAgent {
type Output = String;
fn description(&self) -> &str {
"Benchmark agent"
}
fn output_schema(&self) -> Option<Value> {
None
}
fn name(&self) -> &str {
"bench_agent"
}
fn tools(&self) -> Vec<Box<dyn ToolT>> {
vec![Box::new(BenchTool) as Box<dyn ToolT>]
}
}
impl AgentHooks for BenchAgent {}
#[derive(Debug, Clone)]
struct BenchConfig {
prompt: String,
llm_delay_ms: u64,
}
impl BenchConfig {
fn from_env() -> Self {
Self {
prompt: env_or_default(
"AUTOAGENTS_BENCH_PROMPT",
"Answer with a single number: 2+2".to_string(),
),
llm_delay_ms: env_u64("AUTOAGENTS_BENCH_LLM_DELAY_MS", DEFAULT_LLM_DELAY_MS),
}
}
}
fn env_or_default(key: &str, default: String) -> String {
std::env::var(key).unwrap_or(default)
}
fn env_u64(key: &str, default: u64) -> u64 {
std::env::var(key).ok().map_or(default, |value| {
value
.parse::<u64>()
.unwrap_or_else(|_| panic!("Invalid {key} value (expected u64): {value}"))
})
}
fn build_llm(config: &BenchConfig) -> Arc<dyn LLMProvider> {
Arc::new(MockLlmProvider::new(config.llm_delay_ms))
}
async fn build_direct_agent(
llm: Arc<dyn LLMProvider>,
) -> DirectAgentHandle<ReActAgent<BenchAgent>> {
let sliding_window_memory = Box::new(SlidingWindowMemory::new(8));
let agent = ReActAgent::new(BenchAgent);
AgentBuilder::<_, DirectAgent>::new(agent)
.llm(llm)
.memory(sliding_window_memory)
.build()
.await
.expect("Failed to build agent")
}
async fn build_actor_agent(
llm: Arc<dyn LLMProvider>,
runtime: Arc<SingleThreadedRuntime>,
) -> ActorAgentHandle<ReActAgent<BenchAgent>> {
let sliding_window_memory = Box::new(SlidingWindowMemory::new(8));
let agent = ReActAgent::new(BenchAgent);
AgentBuilder::<_, ActorAgent>::new(agent)
.llm(llm)
.memory(sliding_window_memory)
.runtime(runtime)
.build()
.await
.expect("Failed to build actor agent")
}
type BenchDirectAgent = autoagents_core::agent::BaseAgent<ReActAgent<BenchAgent>, DirectAgent>;
#[derive(Debug, Serialize)]
struct MemorySample {
rss_before_kb: u64,
rss_after_kb: u64,
rss_delta_kb: u64,
}
#[derive(Debug, Serialize)]
struct MemoryReport {
platform: String,
available: bool,
samples: Vec<MemorySample>,
}
static MEMORY_SAMPLES: OnceLock<Mutex<Vec<MemorySample>>> = OnceLock::new();
fn memory_samples() -> &'static Mutex<Vec<MemorySample>> {
MEMORY_SAMPLES.get_or_init(|| Mutex::new(Vec::new()))
}
#[cfg(target_os = "linux")]
fn read_rss_kb() -> Option<u64> {
let status = std::fs::read_to_string("/proc/self/status").ok()?;
for line in status.lines() {
if let Some(rest) = line.strip_prefix("VmRSS:") {
let value = rest.split_whitespace().next()?;
return value.parse::<u64>().ok();
}
}
None
}
#[cfg(not(target_os = "linux"))]
fn read_rss_kb() -> Option<u64> {
None
}
fn write_memory_report() {
let samples = memory_samples()
.lock()
.map(|mut guard| std::mem::take(&mut *guard))
.unwrap_or_default();
let report = MemoryReport {
platform: std::env::consts::OS.to_string(),
available: cfg!(target_os = "linux"),
samples,
};
let out_dir = PathBuf::from("target")
.join("criterion")
.join("agent_runtime");
if let Err(err) = std::fs::create_dir_all(&out_dir) {
eprintln!("Failed to create memory report directory: {err}");
return;
}
let path = out_dir.join("memory.json");
match serde_json::to_string_pretty(&report) {
Ok(json) => {
if let Err(err) = std::fs::write(&path, json) {
eprintln!("Failed to write memory report: {err}");
}
}
Err(err) => eprintln!("Failed to serialize memory report: {err}"),
}
}
struct OverheadReportInputs<'a> {
runtime: &'a TokioRuntime,
llm: &'a Arc<dyn LLMProvider>,
direct_agent: &'a BenchDirectAgent,
actor_ref: &'a autoagents_core::ractor::ActorRef<Task>,
actor_events: &'a mut autoagents_core::utils::BoxEventStream<Event>,
prompt: &'a str,
tools: &'a [Tool],
tool_result_message: &'a ChatMessage,
}
fn run_overhead_report(inputs: OverheadReportInputs<'_>) -> Result<(), String> {
let OverheadReportInputs {
runtime,
llm,
direct_agent,
actor_ref,
actor_events,
prompt,
tools,
tool_result_message,
} = inputs;
let iterations = 10u32;
let message = ChatMessage::user().content(prompt.to_string()).build();
let llm_messages = vec![message.clone()];
let llm_followup = vec![message.clone(), tool_result_message.clone()];
let llm_no_tools_start = Instant::now();
for _ in 0..iterations {
let msg = message.clone();
let _ = runtime.block_on(llm.chat(std::slice::from_ref(&msg), None));
}
let llm_no_tools_elapsed = llm_no_tools_start.elapsed();
let llm_with_tools_start = Instant::now();
for _ in 0..iterations {
let _ = runtime.block_on(llm.chat_with_tools(&llm_messages, Some(tools), None));
let _ = runtime.block_on(llm.chat_with_tools(&llm_followup, Some(tools), None));
}
let llm_with_tools_elapsed = llm_with_tools_start.elapsed();
let direct_start = Instant::now();
for _ in 0..iterations {
let _ = runtime.block_on(direct_agent.run(Task::new(prompt.to_string())));
}
let direct_elapsed = direct_start.elapsed();
let actor_start = Instant::now();
for _ in 0..iterations {
let task = Task::new(prompt.to_string());
let submission_id = task.submission_id;
actor_ref
.cast(task)
.map_err(|e| format!("Failed to send actor task: {e}"))?;
runtime.block_on(await_task_complete(actor_events, submission_id))?;
}
let actor_elapsed = actor_start.elapsed();
let llm_no_tools_avg_ms = llm_no_tools_elapsed.as_secs_f64() * 1000.0 / f64::from(iterations);
let llm_with_tools_avg_ms =
llm_with_tools_elapsed.as_secs_f64() * 1000.0 / f64::from(iterations);
let direct_avg_ms = direct_elapsed.as_secs_f64() * 1000.0 / f64::from(iterations);
let actor_avg_ms = actor_elapsed.as_secs_f64() * 1000.0 / f64::from(iterations);
let tool_cycle_overhead_ms = llm_with_tools_avg_ms - llm_no_tools_avg_ms;
let direct_overhead_ms = direct_avg_ms - llm_with_tools_avg_ms;
let actor_overhead_ms = actor_avg_ms - llm_with_tools_avg_ms;
let direct_overhead_pct = if llm_with_tools_avg_ms > 0.0 {
(direct_overhead_ms / llm_with_tools_avg_ms) * 100.0
} else {
0.0
};
let actor_overhead_pct = if llm_with_tools_avg_ms > 0.0 {
(actor_overhead_ms / llm_with_tools_avg_ms) * 100.0
} else {
0.0
};
eprintln!("Overhead summary ({} iters):", iterations);
eprintln!(" llm_no_tools={:.3} ms", llm_no_tools_avg_ms);
eprintln!(" llm_with_tools={:.3} ms", llm_with_tools_avg_ms);
eprintln!(" tool_cycle_overhead={:.3} ms", tool_cycle_overhead_ms);
eprintln!(
" direct_agent={:.3} ms (overhead +{:.3} ms, {:.1}%)",
direct_avg_ms, direct_overhead_ms, direct_overhead_pct
);
eprintln!(
" actor_agent={:.3} ms (overhead +{:.3} ms, {:.1}%)",
actor_avg_ms, actor_overhead_ms, actor_overhead_pct
);
Ok(())
}
fn bench_agent_runtime(c: &mut Criterion) {
let config = BenchConfig::from_env();
{
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("Failed to build tokio runtime");
let llm = build_llm(&config);
let direct_handle = runtime.block_on(build_direct_agent(llm.clone()));
let DirectAgentHandle {
agent: direct_agent,
rx,
..
} = direct_handle;
runtime.spawn(async move {
let mut rx = rx;
while rx.next().await.is_some() {}
});
let actor_runtime = SingleThreadedRuntime::new(Some(DEFAULT_RUNTIME_BUFFER));
let mut actor_events = runtime
.block_on(actor_runtime.take_event_receiver())
.expect("Failed to acquire actor event receiver");
let actor_runtime_task = {
let actor_runtime = actor_runtime.clone();
runtime.spawn(async move {
let _ = actor_runtime.run().await;
})
};
let actor_handle = runtime.block_on(build_actor_agent(llm.clone(), actor_runtime.clone()));
let actor_ref = actor_handle.actor_ref.clone();
let prompt = config.prompt.clone();
let llm_message = ChatMessage::user().content(prompt.clone()).build();
let tool_defs = bench_tool_defs();
let tool_result_message = bench_tool_result_message();
let llm_messages = vec![llm_message.clone()];
let llm_followup = vec![llm_message.clone(), tool_result_message.clone()];
let _ = runtime.block_on(direct_agent.run(Task::new(prompt.clone())));
let warm_task = Task::new(prompt.clone());
let warm_id = warm_task.submission_id;
actor_ref
.cast(warm_task)
.expect("Failed to dispatch warm-up actor task");
runtime
.block_on(await_task_complete(&mut actor_events, warm_id))
.expect("Warm-up actor task failed");
run_overhead_report(OverheadReportInputs {
runtime: &runtime,
llm: &llm,
direct_agent: &direct_agent,
actor_ref: &actor_ref,
actor_events: &mut actor_events,
prompt: &prompt,
tools: &tool_defs,
tool_result_message: &tool_result_message,
})
.expect("Failed to compute overhead report");
let mut group = c.benchmark_group("agent_runtime");
group.bench_function("llm_no_tools", |b| {
b.iter(|| {
let message = llm_message.clone();
let result = runtime.block_on(llm.chat(std::slice::from_ref(&message), None));
black_box(result).ok();
});
});
group.bench_function("llm_with_tools", |b| {
b.iter(|| {
let result =
runtime.block_on(llm.chat_with_tools(&llm_messages, Some(&tool_defs), None));
let follow_up =
runtime.block_on(llm.chat_with_tools(&llm_followup, Some(&tool_defs), None));
black_box(result).ok();
black_box(follow_up).ok();
});
});
group.bench_function("direct_agent", |b| {
b.iter(|| {
let task = Task::new(prompt.clone());
let result = runtime.block_on(direct_agent.run(task));
black_box(result).ok();
});
});
group.bench_function("actor_agent", |b| {
b.iter(|| {
let task = Task::new(prompt.clone());
let submission_id = task.submission_id;
actor_ref.cast(task).expect("Failed to dispatch actor task");
runtime
.block_on(await_task_complete(&mut actor_events, submission_id))
.expect("Actor task failed");
});
});
group.bench_function("direct_agent_rss_delta", |b| {
b.iter_custom(|iters| {
let mut local_samples = Vec::new();
let start = Instant::now();
for _ in 0..iters {
let rss_before = read_rss_kb();
let task = Task::new(prompt.clone());
let result = runtime.block_on(direct_agent.run(task));
black_box(result).ok();
let rss_after = read_rss_kb();
if let (Some(before), Some(after)) = (rss_before, rss_after) {
local_samples.push(MemorySample {
rss_before_kb: before,
rss_after_kb: after,
rss_delta_kb: after.saturating_sub(before),
});
}
}
if !local_samples.is_empty()
&& let Ok(mut guard) = memory_samples().lock()
{
guard.extend(local_samples);
}
start.elapsed()
});
});
group.finish();
let _ = runtime.block_on(actor_runtime.stop());
actor_runtime_task.abort();
}
write_memory_report();
}
criterion_group! {
name = benches;
config = Criterion::default().configure_from_args();
targets = bench_agent_runtime
}
criterion_main!(benches);