use async_trait::async_trait;
use axum::{
extract::{Path, State},
http::StatusCode,
response::sse::{Event, Sse},
Json,
};
use std::collections::HashMap;
use std::convert::Infallible;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use tokio::sync::broadcast;
use tokio_stream::{wrappers::BroadcastStream, wrappers::IntervalStream, StreamExt};
use crate::event::{AgentEvent, ChannelSink, NullSink};
use crate::message::Role;
use crate::permissions::{LayeredPermissionsConfig, PermissionMode};
use crate::runtime::AgentRuntimeBuilder;
use super::{
build_openapi_spec, ApiError, AppState, CreateSessionRequest, CreateSessionResponse,
ErrorResponse, ListSessionsQuery, RunRequest, RunResponse, SessionDetailResponse, SessionInfo,
SessionMessageRequest, SessionMessageResponse, SessionState, SetGoalRequest, SlashCommandInfo,
SseContentBlock, SseEvent, ToolInfo, UsageInfo,
};
#[cfg_attr(test, mutants::skip)]
pub(super) async fn health() -> &'static str {
"ok"
}
fn record_run_success(metrics: &super::Metrics, steps: usize, usage: &crate::llm::TokenUsage) {
metrics.agent_runs_total.fetch_add(1, Ordering::Relaxed);
metrics.agent_runs_success.fetch_add(1, Ordering::Relaxed);
metrics
.agent_steps_total
.fetch_add(steps as u64, Ordering::Relaxed);
metrics
.tokens_prompt_total
.fetch_add(usage.prompt_tokens as u64, Ordering::Relaxed);
metrics
.tokens_completion_total
.fetch_add(usage.completion_tokens as u64, Ordering::Relaxed);
}
fn record_run_failed(metrics: &super::Metrics) {
metrics.agent_runs_total.fetch_add(1, Ordering::Relaxed);
metrics.agent_runs_failed.fetch_add(1, Ordering::Relaxed);
}
fn map_run_error(e: &crate::error::Error) -> ApiError {
use crate::error::Error;
match e {
Error::Cancelled => ApiError::new(StatusCode::SERVICE_UNAVAILABLE, "agent run cancelled"),
Error::PermissionDenied { .. } | Error::PermissionDeniedLimit { .. } => {
ApiError::forbidden(e.to_string())
}
Error::RateLimited { retry_after_ms, .. } => {
ApiError::new(StatusCode::TOO_MANY_REQUESTS, e.to_string())
.with_retry_after((*retry_after_ms / 1000) as u32)
}
Error::BadToolArgs { .. } => ApiError::bad_request(e.to_string()),
_ => ApiError::internal(e.to_string()),
}
}
#[cfg_attr(test, mutants::skip)]
pub(super) async fn openapi_spec() -> Json<serde_json::Value> {
Json(build_openapi_spec())
}
#[cfg_attr(test, mutants::skip)]
pub(super) async fn list_tools(State(state): State<Arc<AppState>>) -> Json<Vec<ToolInfo>> {
Json(state.tools.clone())
}
pub(super) async fn run_agent(
State(state): State<Arc<AppState>>,
Json(body): Json<RunRequest>,
) -> Result<Json<RunResponse>, ApiError> {
if body.goal.trim().is_empty() {
return Err(ApiError::bad_request("missing or empty 'goal' field"));
}
let _permit = state
.run_semaphore
.clone()
.acquire_owned()
.await
.map_err(|_| {
ApiError::new(
StatusCode::SERVICE_UNAVAILABLE,
"too many concurrent runs, try again later",
)
})?;
let max_steps = body.max_steps.unwrap_or(state.config.max_steps as u32) as usize;
let system_prompt = match body.system_prompt {
Some(s) => s,
None => {
let mut p = state.config.system_prompt.clone();
if let Some(extra) = &body.append_system_prompt {
p.push('\n');
p.push_str(extra);
}
p
}
};
let assembled_system_prompt = crate::assemble_system_prompt(
&system_prompt,
&state.config.workspace,
&state.skills,
state.config.subagent_enabled,
);
let system_prompt = assembled_system_prompt.full;
let prompt_segments = assembled_system_prompt.segments;
let mut tool_registry = state.tool_registry.clone();
if let Some(mode_str) = body.permission_mode.as_deref() {
let perm_mode = parse_permission_mode(mode_str, state.config.allow_bypass_permissions);
tool_registry = tool_registry.with_permissions(LayeredPermissionsConfig {
mode: perm_mode,
layers: Vec::new(),
});
}
let mut runtime = AgentRuntimeBuilder::new()
.llm(state.provider.clone())
.tools(tool_registry)
.system_prompt(system_prompt)
.prompt_segments(prompt_segments)
.max_steps(max_steps)
.build()
.map_err(|e| ApiError::internal(format!("failed to build runtime: {e}")))?;
let outcome = runtime.run(&body.goal).await.map_err(|e| {
record_run_failed(&state.metrics);
map_run_error(&e)
})?;
record_run_success(&state.metrics, outcome.steps, &outcome.total_usage);
let messages: Vec<serde_json::Value> = runtime
.transcript()
.iter()
.filter_map(|msg| serde_json::to_value(msg).ok())
.collect();
let finish_reason = outcome.finish_reason.to_string();
Ok(Json(RunResponse {
status: "success".into(),
finish_reason,
messages,
usage: UsageInfo {
total_steps: outcome.steps as u32,
total_tokens: outcome.total_usage.total_tokens as u64,
},
}))
}
fn parse_permission_mode(s: &str, allow_bypass: bool) -> PermissionMode {
match s.to_ascii_lowercase().as_str() {
"auto" => PermissionMode::Auto,
"strict" => PermissionMode::Strict,
"bypass" | "bypass_permissions" if allow_bypass => PermissionMode::BypassPermissions,
_ => PermissionMode::Default,
}
}
#[cfg_attr(test, mutants::skip)]
fn generate_session_id() -> String {
uuid::Uuid::now_v7().to_string()
}
fn format_timestamp(t: SystemTime) -> String {
let dur = t.duration_since(SystemTime::UNIX_EPOCH).unwrap_or_default();
let secs = dur.as_secs();
let days = secs / 86400;
let remaining = secs % 86400;
let hours = remaining / 3600;
let minutes = (remaining % 3600) / 60;
let seconds = remaining % 60;
let (year, month, day) = crate::session::epoch_day_to_ymd(days as i64);
format!(
"{:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
year, month, day, hours, minutes, seconds
)
}
pub(super) async fn create_session(
State(state): State<Arc<AppState>>,
Json(body): Json<CreateSessionRequest>,
) -> Result<(StatusCode, Json<CreateSessionResponse>), ApiError> {
let id = generate_session_id();
let created_at = format_timestamp(SystemTime::now());
let system_prompt = match body.system_prompt {
Some(s) => s,
None => {
let mut p = state.config.system_prompt.clone();
if let Some(extra) = &body.append_system_prompt {
p.push('\n');
p.push_str(extra);
}
p
}
};
let assembled_system_prompt = crate::assemble_system_prompt(
&system_prompt,
&state.config.workspace,
&state.skills,
state.config.subagent_enabled,
);
let system_prompt = assembled_system_prompt.full;
let prompt_segments = assembled_system_prompt.segments;
let max_steps = body
.max_steps
.map(|n| n as usize)
.unwrap_or(state.config.max_steps);
let mut tool_registry = state.tool_registry.clone();
if let Some(mode_str) = body.permission_mode.as_deref() {
let perm_mode = parse_permission_mode(mode_str, state.config.allow_bypass_permissions);
tool_registry = tool_registry.with_permissions(LayeredPermissionsConfig {
mode: perm_mode,
layers: Vec::new(),
});
}
let mut runtime = AgentRuntimeBuilder::new()
.llm(state.provider.clone())
.tools(tool_registry)
.system_prompt(system_prompt)
.prompt_segments(prompt_segments)
.max_steps(max_steps)
.build()
.map_err(|e| ApiError::internal(format!("failed to build session runtime: {e}")))?;
runtime.set_session_id(&id);
let plan_approval_gate = runtime.plan_approval_gate();
let session = SessionState {
id: id.clone(),
created_at: created_at.clone(),
title: body.session_name,
runtime: Arc::new(tokio::sync::Mutex::new(runtime)),
plan_approval_gate,
interrupt_token: Arc::new(tokio::sync::Mutex::new(None)),
non_system_message_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
last_active_ms: Arc::new(AtomicU64::new(super::now_session_ms())),
prompt_tokens: Arc::new(AtomicU64::new(0)),
completion_tokens: Arc::new(AtomicU64::new(0)),
};
state.sessions.write().await.insert(id.clone(), session);
state
.metrics
.sessions_active
.fetch_add(1, Ordering::Relaxed);
tracing::info!(session_id = %id, "session created");
Ok((
StatusCode::CREATED,
Json(CreateSessionResponse { id, created_at }),
))
}
#[derive(serde::Serialize)]
pub(super) struct SessionList {
pub total: usize,
pub sessions: Vec<SessionInfo>,
}
pub(super) async fn list_sessions(
State(state): State<Arc<AppState>>,
axum::extract::Query(params): axum::extract::Query<ListSessionsQuery>,
) -> Json<SessionList> {
let sessions = state.sessions.read().await;
let mut infos = Vec::with_capacity(sessions.len());
for s in sessions.values() {
let message_count = s
.non_system_message_count
.load(std::sync::atomic::Ordering::Relaxed);
infos.push(SessionInfo {
id: s.id.clone(),
created_at: s.created_at.clone(),
message_count,
title: s.title.clone(),
});
}
infos.sort_by(|a, b| a.created_at.cmp(&b.created_at).then(a.id.cmp(&b.id)));
let total = infos.len();
let offset = params.offset.unwrap_or(0);
let page: Vec<SessionInfo> = infos
.into_iter()
.skip(offset)
.take(params.limit.unwrap_or(usize::MAX))
.collect();
Json(SessionList {
total,
sessions: page,
})
}
pub(super) async fn get_session(
State(state): State<Arc<AppState>>,
Path(id): Path<String>,
) -> Result<Json<SessionDetailResponse>, ApiError> {
let sessions = state.sessions.read().await;
let session = sessions
.get(&id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
let pending_plan = session
.plan_approval_gate
.pending_plan
.read()
.ok()
.and_then(|g| g.clone());
let status = if pending_plan.is_some() {
"plan_pending_approval".to_string()
} else {
"idle".to_string()
};
let (messages, todos, goal) = match session.runtime.try_lock() {
Ok(runtime) => {
let msgs = runtime
.transcript()
.iter()
.filter_map(|msg| serde_json::to_value(msg).ok())
.collect();
let todos = runtime.current_todos();
let goal = runtime.current_goal();
(msgs, todos, goal)
}
Err(_) => (vec![], vec![], None),
};
let (first_prompt, last_prompt) = {
let user_msgs: Vec<String> = messages
.iter()
.filter_map(|m| {
if m.get("role")?.as_str()? == "user" {
m.get("content")?.as_str().map(|s| s.to_string())
} else {
None
}
})
.collect();
let first = user_msgs.first().cloned();
let last = user_msgs.last().cloned();
(first, last)
};
let prompt_tokens = session.prompt_tokens.load(Ordering::Relaxed);
let completion_tokens = session.completion_tokens.load(Ordering::Relaxed);
Ok(Json(SessionDetailResponse {
id: session.id.clone(),
created_at: session.created_at.clone(),
title: session.title.clone(),
messages,
todos,
status,
pending_plan,
goal,
first_prompt,
last_prompt,
prompt_tokens,
completion_tokens,
}))
}
pub(super) async fn delete_session(
State(state): State<Arc<AppState>>,
Path(id): Path<String>,
) -> Result<StatusCode, ApiError> {
let session_runtime = {
let sessions = state.sessions.read().await;
sessions.get(&id).map(|s| s.runtime.clone())
};
if let Some(runtime) = session_runtime {
let mut rt = runtime.lock().await;
rt.close(None).await;
drop(rt);
state.sessions.write().await.remove(&id);
state
.metrics
.sessions_active
.fetch_sub(1, Ordering::Relaxed);
state.event_channels.write().await.remove(&id);
tracing::info!(session_id = %id, "session deleted");
Ok(StatusCode::NO_CONTENT)
} else {
Err(ApiError::not_found("session not found"))
}
}
#[derive(serde::Deserialize, Debug)]
pub(super) struct PatchSessionRequest {
title: Option<String>,
}
pub(super) async fn patch_session(
State(state): State<Arc<AppState>>,
Path(id): Path<String>,
Json(body): Json<PatchSessionRequest>,
) -> Result<Json<SessionInfo>, ApiError> {
let mut sessions = state.sessions.write().await;
let session = sessions
.get_mut(&id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
if let Some(title) = body.title {
session.title = if title.is_empty() { None } else { Some(title) };
}
Ok(Json(SessionInfo {
id: session.id.clone(),
created_at: session.created_at.clone(),
message_count: session.non_system_message_count.load(Ordering::Relaxed),
title: session.title.clone(),
}))
}
#[derive(serde::Serialize)]
pub(super) struct ForkSessionResponse {
id: String,
created_at: String,
message_count: usize,
}
pub(super) async fn fork_session(
State(state): State<Arc<AppState>>,
Path(id): Path<String>,
) -> Result<(StatusCode, Json<ForkSessionResponse>), ApiError> {
let transcript_snapshot = {
let sessions = state.sessions.read().await;
let src = sessions
.get(&id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
let rt = src
.runtime
.try_lock()
.map_err(|_| ApiError::conflict("session is busy"))?;
rt.transcript().to_vec()
};
let new_id = generate_session_id();
let created_at = format_timestamp(SystemTime::now());
let base_system_prompt = state.config.system_prompt.clone();
let assembled_system_prompt = crate::assemble_system_prompt(
&base_system_prompt,
&state.config.workspace,
&state.skills,
state.config.subagent_enabled,
);
let system_prompt = assembled_system_prompt.full;
let prompt_segments = assembled_system_prompt.segments;
let mut runtime = AgentRuntimeBuilder::new()
.llm(state.provider.clone())
.tools(state.tool_registry.clone())
.system_prompt(system_prompt)
.prompt_segments(prompt_segments)
.max_steps(state.config.max_steps)
.build()
.map_err(|_| ApiError::internal("failed to build forked session runtime"))?;
let non_system_count = transcript_snapshot
.iter()
.filter(|m| m.role != crate::message::Role::System)
.count();
runtime.set_transcript(transcript_snapshot);
let plan_approval_gate = runtime.plan_approval_gate();
let session = SessionState {
id: new_id.clone(),
created_at: created_at.clone(),
title: None,
runtime: Arc::new(tokio::sync::Mutex::new(runtime)),
plan_approval_gate,
interrupt_token: Arc::new(tokio::sync::Mutex::new(None)),
non_system_message_count: Arc::new(std::sync::atomic::AtomicUsize::new(non_system_count)),
last_active_ms: Arc::new(AtomicU64::new(super::now_session_ms())),
prompt_tokens: Arc::new(AtomicU64::new(0)),
completion_tokens: Arc::new(AtomicU64::new(0)),
};
state.sessions.write().await.insert(new_id.clone(), session);
state
.metrics
.sessions_active
.fetch_add(1, Ordering::Relaxed);
Ok((
StatusCode::CREATED,
Json(ForkSessionResponse {
id: new_id,
created_at,
message_count: non_system_count,
}),
))
}
#[derive(serde::Deserialize)]
pub(super) struct PlanConfirmRequest {
edits: Option<String>,
}
#[derive(serde::Deserialize)]
pub(super) struct PlanRejectRequest {
reason: Option<String>,
}
pub(super) async fn session_plan_confirm(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
Json(body): Json<PlanConfirmRequest>,
) -> Result<Json<serde_json::Value>, ApiError> {
let sessions = state.sessions.read().await;
let session = sessions
.get(&session_id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
let pending = session
.plan_approval_gate
.pending_plan
.read()
.ok()
.and_then(|g| g.clone());
if pending.is_none() {
return Err(ApiError::conflict("session is not awaiting plan approval"));
}
if let Some(edited) = body.edits {
if let Ok(mut w) = session.plan_approval_gate.pending_plan.write() {
*w = Some(edited);
}
}
session.plan_approval_gate.approve();
Ok(Json(serde_json::json!({
"status": "approved",
"session_id": session_id
})))
}
pub(super) async fn session_plan_reject(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
Json(body): Json<PlanRejectRequest>,
) -> Result<Json<serde_json::Value>, ApiError> {
let sessions = state.sessions.read().await;
let session = sessions
.get(&session_id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
let pending = session
.plan_approval_gate
.pending_plan
.read()
.ok()
.and_then(|g| g.clone());
if pending.is_none() {
return Err(ApiError::conflict("session is not awaiting plan approval"));
}
let reason = body.reason.unwrap_or_default();
session.plan_approval_gate.reject(&reason);
Ok(Json(serde_json::json!({
"status": "rejected",
"session_id": session_id
})))
}
pub(super) async fn session_set_goal(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
Json(body): Json<SetGoalRequest>,
) -> Result<Json<serde_json::Value>, ApiError> {
let runtime_arc = {
let sessions = state.sessions.read().await;
let session = sessions
.get(&session_id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
session.runtime.clone()
};
let condition = body.condition.clone();
let max_turns = body.max_turns.unwrap_or(20);
match runtime_arc.try_lock() {
Ok(runtime) => {
runtime.set_goal(condition, max_turns).await;
}
Err(_) => {
return Err(ApiError::conflict("session runtime is busy"));
}
}
Ok(Json(serde_json::json!({
"status": "pursuing",
"session_id": session_id
})))
}
pub(super) async fn session_clear_goal(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
) -> Result<Json<serde_json::Value>, ApiError> {
let runtime_arc = {
let sessions = state.sessions.read().await;
let session = sessions
.get(&session_id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
session.runtime.clone()
};
let lock_result = runtime_arc.try_lock();
match lock_result {
Ok(runtime) => {
runtime.clear_goal().await;
drop(runtime);
Ok(Json(serde_json::json!({
"status": "cleared",
"session_id": session_id
})))
}
Err(_) => {
if runtime_goal_state_clear(&runtime_arc).await {
return Ok(Json(serde_json::json!({
"status": "cleared",
"session_id": session_id
})));
}
Err(ApiError::conflict(
"session runtime is busy; goal not cleared — retry after the current turn completes",
)
.with_retry_after(5))
}
}
}
async fn runtime_goal_state_clear(
runtime: &Arc<tokio::sync::Mutex<crate::runtime::AgentRuntime>>,
) -> bool {
for _ in 0..10u8 {
if let Ok(rt) = runtime.try_lock() {
rt.clear_goal().await;
return true;
}
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
false
}
pub(super) async fn session_interrupt(
State(state): State<Arc<AppState>>,
Path(session_id): Path<String>,
) -> Result<Json<serde_json::Value>, ApiError> {
let token_arc = {
let sessions = state.sessions.read().await;
let session = sessions
.get(&session_id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
session.interrupt_token.clone()
};
let token_opt = token_arc.lock().await.clone();
if let Some(token) = token_opt {
token.cancel();
}
Ok(Json(serde_json::json!({
"status": "interrupted",
"session_id": session_id
})))
}
#[cfg_attr(test, mutants::skip)]
pub(super) async fn list_slash_commands(
State(state): State<Arc<AppState>>,
) -> Json<Vec<SlashCommandInfo>> {
Json((*state.slash_commands).clone())
}
pub(super) async fn send_session_message(
State(state): State<Arc<AppState>>,
Path(id): Path<String>,
Json(body): Json<SessionMessageRequest>,
) -> Result<Json<SessionMessageResponse>, ApiError> {
if body.content.trim().is_empty() {
return Err(ApiError::bad_request("missing or empty 'content' field"));
}
tracing::debug!(session_id = %id, content_len = body.content.len(), "session message received");
let (runtime_arc, interrupt_token_arc, msg_count_arc, prompt_tokens_arc, completion_tokens_arc) = {
let sessions = state.sessions.read().await;
let session = sessions
.get(&id)
.ok_or_else(|| ApiError::not_found("session not found"))?;
session
.last_active_ms
.store(super::now_session_ms(), Ordering::Relaxed);
(
session.runtime.clone(),
session.interrupt_token.clone(),
session.non_system_message_count.clone(),
session.prompt_tokens.clone(),
session.completion_tokens.clone(),
)
};
let broadcast_tx = {
let mut channels = state.event_channels.write().await;
let tx = channels.entry(id.clone()).or_insert_with(|| {
let (tx, _) = broadcast::channel(64);
tx
});
tx.clone()
};
let _permit = state
.run_semaphore
.clone()
.acquire_owned()
.await
.map_err(|_| {
ApiError::new(
StatusCode::SERVICE_UNAVAILABLE,
"too many concurrent runs, try again later",
)
})?;
let mut runtime = runtime_arc.lock().await;
let interrupt_token = tokio_util::sync::CancellationToken::new();
{
let mut stored = interrupt_token_arc.lock().await;
*stored = Some(interrupt_token.clone());
}
runtime.set_interrupt_token(interrupt_token);
let (sink, mut event_rx) = ChannelSink::new();
runtime.set_event_sink(Arc::new(sink));
let initial_count = msg_count_arc.load(std::sync::atomic::Ordering::Relaxed);
let count_arc = msg_count_arc.clone();
let forward_handle = tokio::spawn(async move {
let mut tool_start_times: HashMap<String, std::time::Instant> = HashMap::new();
let mut count: usize = initial_count;
while let Some(ref agent_event) = event_rx.recv().await {
match agent_event {
AgentEvent::MessageAppended { message, .. }
| AgentEvent::MessageAppendedWithAudit { message, .. }
if message.role != Role::System =>
{
count += 1;
count_arc.store(count, std::sync::atomic::Ordering::Relaxed);
}
_ => {}
}
if let AgentEvent::ToolCall { id, .. } = agent_event {
tool_start_times.insert(id.clone(), std::time::Instant::now());
}
if let Some(sse_event) = map_agent_event(agent_event) {
let _ = broadcast_tx.send(sse_event);
}
if let AgentEvent::ToolResult { id, name, .. } = agent_event {
let elapsed_ms = tool_start_times
.remove(id)
.map(|start| start.elapsed().as_millis() as u64)
.unwrap_or(0);
let _ = broadcast_tx.send(SseEvent::ToolProgress {
tool_use_id: id.clone(),
tool_name: name.clone(),
elapsed_ms,
});
}
}
});
let run_result = runtime.enqueue(&body.content).await.map(|opt| {
opt.unwrap_or_else(|| crate::runtime::RuntimeOutcome {
final_text: None,
finish_reason: crate::agent::FinishReason::NoMoreToolCalls,
total_usage: crate::TokenUsage::default(),
steps: 0,
llm_latency_ms: 0,
checkpoint_id: None,
})
});
{
let mut stored = interrupt_token_arc.lock().await;
*stored = None;
}
runtime.set_event_sink(Arc::new(NullSink));
let _ = forward_handle.await;
let outcome = run_result.map_err(|e| {
record_run_failed(&state.metrics);
map_run_error(&e)
})?;
prompt_tokens_arc.fetch_add(outcome.total_usage.prompt_tokens as u64, Ordering::Relaxed);
completion_tokens_arc.fetch_add(
outcome.total_usage.completion_tokens as u64,
Ordering::Relaxed,
);
record_run_success(&state.metrics, outcome.steps, &outcome.total_usage);
let last_assistant = runtime
.transcript()
.iter()
.rev()
.find(|m| m.role == crate::message::Role::Assistant)
.map(|m| m.content.clone())
.unwrap_or_default();
Ok(Json(SessionMessageResponse {
role: "assistant".into(),
content: last_assistant,
}))
}
pub(super) async fn session_events(
State(state): State<Arc<AppState>>,
Path(id): Path<String>,
) -> Result<Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>, ApiError> {
{
let sessions = state.sessions.read().await;
if !sessions.contains_key(&id) {
return Err(ApiError::not_found("session not found"));
}
}
let rx = {
let mut channels = state.event_channels.write().await;
let tx = channels.entry(id.clone()).or_insert_with(|| {
let (tx, _) = broadcast::channel(64);
tx
});
tx.subscribe()
};
let agent_stream = BroadcastStream::new(rx).filter_map(|result| match result {
Ok(sse_event) => {
let event_type = match &sse_event {
SseEvent::Message { .. } => "message",
SseEvent::PartialMessage { .. } => "partial_message",
SseEvent::ToolCall { .. } => "tool_call",
SseEvent::ToolResult { .. } => "tool_result",
SseEvent::Done { .. } => "done",
SseEvent::Error { .. } => "error",
SseEvent::PlanProposed { .. } => "plan_proposed",
SseEvent::GoalContinuing { .. } => "goal_continuing",
SseEvent::GoalAchieved { .. } => "goal_achieved",
SseEvent::ToolProgress { .. } => "tool_progress",
};
let data = serde_json::to_string(&sse_event).unwrap_or_default();
Some(Ok::<Event, Infallible>(
Event::default().event(event_type).data(data),
))
}
Err(_) => None,
});
let heartbeat_stream = IntervalStream::new(tokio::time::interval(Duration::from_secs(30)))
.map(|_| Ok::<Event, Infallible>(Event::default().comment("heartbeat")));
let combined = agent_stream
.merge(heartbeat_stream)
.timeout(Duration::from_secs(3600))
.filter_map(|r| r.ok());
Ok(Sse::new(combined))
}
pub fn map_agent_event(event: &AgentEvent) -> Option<SseEvent> {
match event {
AgentEvent::PartialToken { text, step } => Some(SseEvent::PartialMessage {
text: text.clone(),
step: *step,
}),
AgentEvent::MessageAppended { message, .. }
| AgentEvent::MessageAppendedWithAudit { message, .. } => {
sse_message_from_canonical(message)
}
AgentEvent::ToolCall { name, step, .. } => Some(SseEvent::ToolCall {
name: name.clone(),
step: *step,
}),
AgentEvent::ToolResult { name, is_error, .. } => {
let success = !is_error;
Some(SseEvent::ToolResult {
name: name.clone(),
success,
})
}
AgentEvent::TurnFinished { reason, steps } => Some(SseEvent::Done {
finish_reason: reason.clone(),
total_steps: *steps,
}),
AgentEvent::PlanProposed { plan_text, .. } => Some(SseEvent::PlanProposed {
plan: plan_text.clone(),
}),
AgentEvent::GoalContinuing { reason, turns } => Some(SseEvent::GoalContinuing {
reason: reason.clone(),
turns: *turns,
}),
AgentEvent::GoalAchieved { condition, turns } => Some(SseEvent::GoalAchieved {
condition: condition.clone(),
turns: *turns,
}),
_ => None,
}
}
fn sse_message_from_canonical(msg: &crate::message::Message) -> Option<SseEvent> {
use crate::message::Role;
let role = match msg.role {
Role::Assistant => "assistant",
Role::User => "user",
Role::System | Role::Tool => return None,
};
let mut content: Vec<SseContentBlock> = Vec::new();
if !msg.content.is_empty() {
content.push(SseContentBlock::Text {
text: msg.content.clone(),
});
}
for tc in &msg.tool_calls {
content.push(SseContentBlock::ToolUse {
id: tc.id.clone(),
name: tc.name.clone(),
input: tc.arguments.clone(),
});
}
if content.is_empty() {
return None;
}
Some(SseEvent::Message {
role: role.into(),
content,
})
}
#[derive(Default)]
struct AguiConverter {
open_message_id: Option<String>,
last_assistant_message_id: Option<String>,
}
impl AguiConverter {
fn new() -> Self {
Self::default()
}
fn convert(&mut self, ev: &AgentEvent) -> Vec<agui_protocol::Event> {
use agui_protocol as ag;
let mut out = Vec::new();
match ev {
AgentEvent::AssistantText { text, .. } => {
if let Some(id) = self.open_message_id.take() {
out.push(ag::Event::TextMessageEnd(ag::TextMessageEnd {
message_id: id,
base: ag::BaseEvent::default(),
}));
}
let id = uuid::Uuid::new_v4().to_string();
out.push(ag::Event::TextMessageStart(ag::TextMessageStart {
message_id: id.clone(),
role: Some("assistant".into()),
base: ag::BaseEvent::default(),
}));
out.push(ag::Event::TextMessageContent(ag::TextMessageContent {
message_id: id.clone(),
delta: text.clone(),
base: ag::BaseEvent::default(),
}));
out.push(ag::Event::TextMessageEnd(ag::TextMessageEnd {
message_id: id.clone(),
base: ag::BaseEvent::default(),
}));
self.last_assistant_message_id = Some(id);
self.open_message_id = None;
}
AgentEvent::PartialToken { text, .. } => {
let id = if let Some(id) = self.open_message_id.clone() {
id
} else {
let id = uuid::Uuid::new_v4().to_string();
out.push(ag::Event::TextMessageStart(ag::TextMessageStart {
message_id: id.clone(),
role: Some("assistant".into()),
base: ag::BaseEvent::default(),
}));
self.open_message_id = Some(id.clone());
self.last_assistant_message_id = Some(id.clone());
id
};
out.push(ag::Event::TextMessageContent(ag::TextMessageContent {
message_id: id,
delta: text.clone(),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::ToolCall {
id,
name,
arguments,
..
} => {
if let Some(open) = self.open_message_id.take() {
out.push(ag::Event::TextMessageEnd(ag::TextMessageEnd {
message_id: open,
base: ag::BaseEvent::default(),
}));
}
out.push(ag::Event::ToolCallStart(ag::ToolCallStart {
tool_call_id: id.clone(),
tool_call_name: name.clone(),
parent_message_id: self.last_assistant_message_id.clone(),
base: ag::BaseEvent::default(),
}));
out.push(ag::Event::ToolCallArgs(ag::ToolCallArgs {
tool_call_id: id.clone(),
delta: arguments.clone(),
base: ag::BaseEvent::default(),
}));
out.push(ag::Event::ToolCallEnd(ag::ToolCallEnd {
tool_call_id: id.clone(),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::ToolResult { id, output, .. } => {
let message_id = self
.last_assistant_message_id
.clone()
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
out.push(ag::Event::ToolCallResult(ag::ToolCallResult {
tool_call_id: id.clone(),
message_id,
content: output.clone(),
role: Some("tool".into()),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::TurnFinished { .. } => {
if let Some(open) = self.open_message_id.take() {
out.push(ag::Event::TextMessageEnd(ag::TextMessageEnd {
message_id: open,
base: ag::BaseEvent::default(),
}));
}
}
AgentEvent::HookStarted {
hook_event,
hook_name,
status_message,
..
} => {
out.push(ag::Event::Custom(ag::Custom {
name: "agui-tui/hook_started".into(),
value: serde_json::json!({
"hookEvent": hook_event,
"hookName": hook_name,
"statusMessage": status_message,
}),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::HookProgress {
hook_event,
hook_name,
last_line,
..
} => {
out.push(ag::Event::Custom(ag::Custom {
name: "agui-tui/hook_progress".into(),
value: serde_json::json!({
"hookEvent": hook_event,
"hookName": hook_name,
"lastLine": last_line,
}),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::HookFinished {
hook_event,
hook_name,
outcome,
duration_ms,
..
} => {
out.push(ag::Event::Custom(ag::Custom {
name: "agui-tui/hook_finished".into(),
value: serde_json::json!({
"hookEvent": hook_event,
"hookName": hook_name,
"outcome": outcome,
"durationMs": duration_ms,
}),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::HookSystemMessage { text, .. } => {
out.push(ag::Event::Custom(ag::Custom {
name: "agui-tui/hook_system_message".into(),
value: serde_json::json!({ "text": text }),
base: ag::BaseEvent::default(),
}));
}
AgentEvent::TodoUpdated { todos, .. } => {
out.push(ag::Event::Custom(ag::Custom {
name: "agui-tui/todo_updated".into(),
value: serde_json::json!({ "todos": todos }),
base: ag::BaseEvent::default(),
}));
}
_ => {}
}
out
}
}
#[cfg(test)]
fn agui_events_for(ev: &AgentEvent) -> Vec<agui_protocol::Event> {
AguiConverter::new().convert(ev)
}
pub(super) async fn agui_run(
State(state): State<Arc<AppState>>,
Json(body): Json<serde_json::Value>,
) -> Result<
Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>,
(StatusCode, Json<ErrorResponse>),
> {
use agui_protocol as ag;
let input: ag::RunAgentInput = serde_json::from_value(body).map_err(|e| {
(
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
status: "error".into(),
error: format!("invalid AG-UI RunAgentInput: {e}"),
}),
)
})?;
let resume_items: Vec<ag::Resume> = input.resume.unwrap_or_default();
let goal = if resume_items.is_empty() {
input
.messages
.iter()
.rev()
.find(|m| m.role == "user")
.and_then(|m| m.content.clone())
.or_else(|| input.context.first().map(|c| c.value.clone()))
.filter(|s| !s.trim().is_empty())
.ok_or_else(|| {
(
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
status: "error".into(),
error: "RunAgentInput must contain at least one user \
message or a non-empty context item"
.into(),
}),
)
})?
} else {
"[frontend tool result received] 客户端工具结果已注入对话,请基于该结果继续回答用户最初的问题。".to_string()
};
let mut seed_transcript: Option<Vec<crate::message::Message>> = None;
if !resume_items.is_empty() {
let session_dir =
agui_session_dir(&state.config.workspace, &input.thread_id).ok_or_else(|| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
status: "error".into(),
error: "cannot resolve session directory for resume".into(),
}),
)
})?;
if !session_dir.join("transcript.jsonl").is_file() {
return Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
status: "error".into(),
error: format!(
"no prior run found for thread '{}'; cannot resume",
input.thread_id
),
}),
));
}
let open_interrupts = load_open_interrupts(&session_dir);
if open_interrupts.is_empty() {
return Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
status: "error".into(),
error: format!(
"thread '{}' has no open interrupts; nothing to resume",
input.thread_id
),
}),
));
}
let resume_ids: std::collections::HashSet<String> = resume_items
.iter()
.map(|r| r.interrupt_id.clone())
.collect();
for open_int in &open_interrupts {
if !resume_ids.contains(&open_int.interrupt_id) {
return Err((
StatusCode::BAD_REQUEST,
Json(ErrorResponse {
status: "error".into(),
error: format!(
"resume must cover all open interrupts; missing '{}'",
open_int.interrupt_id
),
}),
));
}
}
let transcript_path = session_dir.join("transcript.jsonl");
let loaded_messages: Vec<crate::message::Message> =
std::fs::read_to_string(&transcript_path)
.map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
status: "error".into(),
error: format!("failed to load session transcript: {e}"),
}),
)
})?
.lines()
.filter(|line| !line.trim().is_empty())
.filter_map(|line| serde_json::from_str(line).ok())
.collect();
let resume_by_id: std::collections::HashMap<&str, &ag::Resume> = resume_items
.iter()
.map(|r| (r.interrupt_id.as_str(), r))
.collect();
let mut modified = loaded_messages;
for open_int in &open_interrupts {
let Some(resume) = resume_by_id.get(open_int.interrupt_id.as_str()) else {
continue;
};
let tool_call_id = &open_int.tool_call_id;
if resume.status == ag::ResumeStatus::Cancelled {
let sentinel = crate::message::Message::tool_result(
tool_call_id,
"[interrupt cancelled by user]",
);
modified.push(sentinel);
} else if let Some(ref payload) = resume.payload {
let payload_str = serde_json::to_string(payload).unwrap_or_default();
let replaced = modified.iter_mut().any(|msg| {
if msg.tool_call_id.as_deref() == Some(tool_call_id.as_str()) {
msg.content = payload_str.clone();
true
} else {
false
}
});
if !replaced {
modified.push(crate::message::Message::tool_result(
tool_call_id,
&payload_str,
));
}
}
}
clear_open_interrupts(&session_dir);
seed_transcript = Some(modified);
}
if resume_items.is_empty() {
if let Some(session_dir) = agui_session_dir(&state.config.workspace, &input.thread_id) {
let open = load_open_interrupts(&session_dir);
if !open.is_empty() {
return Err((
StatusCode::CONFLICT,
Json(ErrorResponse {
status: "error".into(),
error: format!(
"thread '{}' has {} open interrupt(s); \
must provide resume to continue",
input.thread_id,
open.len()
),
}),
));
}
}
}
let _permit = state
.run_semaphore
.clone()
.try_acquire_owned()
.map_err(|_| {
(
StatusCode::SERVICE_UNAVAILABLE,
Json(ErrorResponse {
status: "error".into(),
error: "too many concurrent runs, try again later".into(),
}),
)
})?;
let assembled_system_prompt = crate::assemble_system_prompt(
&state.config.system_prompt,
&state.config.workspace,
&state.skills,
state.config.subagent_enabled,
);
let system_prompt = assembled_system_prompt.full;
let prompt_segments = assembled_system_prompt.segments;
let interrupt_hook: Option<Arc<TestInterruptHook>> =
input.interrupt_before.as_ref().map(|names| {
Arc::new(TestInterruptHook {
interrupt_before: names.clone(),
interrupted_tool_call_id: std::sync::Mutex::new(None),
interrupted_tool_name: std::sync::Mutex::new(None),
interrupted_arguments: std::sync::Mutex::new(None),
})
});
let agui_tools = input.tools.clone();
let client_tool_names: std::collections::HashSet<String> =
agui_tools.iter().map(|t| t.name.clone()).collect();
let mut tool_registry = state.tool_registry.clone();
for t in &agui_tools {
tool_registry = tool_registry.register(Arc::new(ClientToolStub {
name: t.name.clone(),
description: t.description.clone(),
parameters: t.parameters.clone(),
}));
}
let client_hook: Option<Arc<ClientToolHook>> = if client_tool_names.is_empty() {
None
} else {
Some(Arc::new(ClientToolHook {
names: client_tool_names,
denied: std::sync::Mutex::new(None),
}))
};
if let Some(ref client_hook_ref) = client_hook {
let hook: Arc<ClientToolHook> = client_hook_ref.clone();
tool_registry.set_permission_hook(hook);
}
if client_hook.is_none() {
if let Some(ref test_hook_ref) = interrupt_hook {
let hook: Arc<TestInterruptHook> = test_hook_ref.clone();
tool_registry.set_permission_hook(hook);
}
}
let mut runtime_builder = AgentRuntimeBuilder::new()
.llm(state.provider.clone())
.tools(tool_registry)
.system_prompt(system_prompt)
.prompt_segments(prompt_segments)
.max_steps(state.config.max_steps);
if let Some(seed) = seed_transcript {
runtime_builder = runtime_builder.seed_transcript(seed);
}
let mut runtime = runtime_builder.build().map_err(|e| {
(
StatusCode::INTERNAL_SERVER_ERROR,
Json(ErrorResponse {
status: "error".into(),
error: format!("failed to build runtime: {e}"),
}),
)
})?;
if let Some(ref hook) = client_hook {
runtime.set_permission_hook(hook.clone());
} else if let Some(ref hook) = interrupt_hook {
runtime.set_permission_hook(hook.clone());
}
if let Ok(repo) = crate::ShadowRepo::open(&state.config.workspace) {
let session_id = sanitize_thread_id_for_session(&input.thread_id);
if let Ok(session_dir) = crate::user_sessions_dir(&state.config.workspace) {
let log_dir = session_dir.join(format!("agui-{session_id}"));
let _ = std::fs::create_dir_all(&log_dir);
let log_path = log_dir.join("checkpoints.jsonl");
let touched = runtime.kernel().tools().touched_files();
if let Err(e) =
runtime.enable_checkpoints(Arc::new(repo), session_id, log_path, touched)
{
tracing::warn!("agui: enable_checkpoints failed, continuing without: {e}");
}
}
} else {
tracing::debug!("agui: shadow git unavailable, no per-turn checkpoints");
}
let (sink, mut event_rx) = ChannelSink::new();
runtime.set_event_sink(Arc::new(sink));
let (sse_tx, sse_rx) = tokio::sync::mpsc::unbounded_channel::<ag::Event>();
let thread_id = input.thread_id.clone();
let run_id = input.run_id.clone();
let _ = sse_tx.send(ag::Event::RunStarted(ag::RunStarted {
thread_id: thread_id.clone(),
run_id: run_id.clone(),
base: ag::BaseEvent::default(),
}));
let conv_tx = sse_tx.clone();
let converter_handle = tokio::spawn(async move {
let mut conv = AguiConverter::new();
while let Some(agent_event) = event_rx.recv().await {
for ev in conv.convert(&agent_event) {
if conv_tx.send(ev).is_err() {
return;
}
}
}
});
let metrics = state.metrics.clone();
let drv_thread = thread_id.clone();
let drv_run = run_id.clone();
let drv_interrupt_hook = interrupt_hook.clone();
let drv_client_hook = client_hook.clone();
let drv_client_tools = agui_tools.clone();
let drv_workspace = state.config.workspace.clone();
let driver_handle = tokio::spawn(async move {
let outcome = runtime.run(&goal).await;
let find_denied_tool_call = |transcript: &[crate::message::Message], marker: &str| {
transcript
.iter()
.rev()
.find(|msg| msg.role == crate::message::Role::Tool && msg.content.contains(marker))
.and_then(|msg| msg.tool_call_id.clone())
};
let client_denied: Option<(String, String)> = drv_client_hook.as_ref().and_then(|h| {
let guard = h.denied.lock().unwrap_or_else(|e| e.into_inner());
guard.clone()
});
let test_was_interrupted = drv_interrupt_hook
.as_ref()
.and_then(|hook| {
hook.interrupted_tool_name
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
})
.is_some();
runtime.set_event_sink(Arc::new(NullSink));
let (checkpoint_id, finished_turn): (Option<String>, Option<usize>) = match &outcome {
Ok(o) => (
o.checkpoint_id.as_ref().map(|c| c.0.clone()),
runtime.turn_index().checked_sub(1),
),
Err(_) => (None, None),
};
match &outcome {
Ok(o) => record_run_success(&metrics, o.steps, &o.total_usage),
Err(_) => record_run_failed(&metrics),
}
if let Some(session_dir) = agui_session_dir(&drv_workspace, &drv_thread) {
let _ = std::fs::create_dir_all(&session_dir);
let mut lines = String::new();
for m in runtime.transcript() {
if let Ok(v) = serde_json::to_string(m) {
lines.push_str(&v);
lines.push('\n');
}
}
let path = session_dir.join("transcript.jsonl");
if let Err(e) = std::fs::write(&path, lines) {
tracing::warn!("agui: transcript persist failed: {e}");
}
}
let interrupt_details: Option<AguiInterruptDetail> =
if let Some((name, args)) = client_denied {
let transcript = runtime.transcript();
find_denied_tool_call(transcript, CLIENT_TOOL_DENY_PREFIX).map(|tc_id| {
let parameters = drv_client_tools
.iter()
.find(|t| t.name == name)
.map(|t| t.parameters.clone())
.unwrap_or_else(|| serde_json::json!({"type": "object"}));
AguiInterruptDetail::Client {
tool_call_id: tc_id,
tool_name: name,
parameters,
args: serde_json::Value::String(args),
}
})
} else if test_was_interrupted {
let transcript = runtime.transcript();
let denied_tool_name = drv_interrupt_hook.as_ref().and_then(|h| {
h.interrupted_tool_name
.lock()
.unwrap_or_else(|e| e.into_inner())
.clone()
});
find_denied_tool_call(transcript, "test interrupt trigger").map(|tc_id| {
AguiInterruptDetail::Test {
tool_call_id: tc_id,
tool_name: denied_tool_name.unwrap_or_else(|| "unknown".into()),
}
})
} else {
None
};
let _ = converter_handle.await;
if let (Some(cp), Some(turn)) = (checkpoint_id, finished_turn) {
let _ = sse_tx.send(ag::Event::Custom(ag::Custom {
name: "agui-tui/checkpoint_post".into(),
value: serde_json::json!({
"turn": turn,
"postId": cp,
}),
base: ag::BaseEvent::default(),
}));
}
if let Some(detail) = interrupt_details {
let (tc_id, _tc_name, response_schema, metadata, message) = match detail {
AguiInterruptDetail::Test {
tool_call_id,
tool_name,
} => (
tool_call_id,
tool_name.clone(),
serde_json::json!({
"type": "object",
"properties": { "approved": { "type": "boolean" } }
}),
serde_json::json!({ "testTrigger": true, "toolName": tool_name }),
format!("Test interrupt: tool '{tool_name}' needs user input to proceed"),
),
AguiInterruptDetail::Client {
tool_call_id,
tool_name,
parameters,
args,
} => (
tool_call_id,
tool_name.clone(),
parameters,
serde_json::json!({ "frontendTool": true, "toolName": tool_name, "args": args }),
format!("Frontend tool '{tool_name}' needs client execution to proceed"),
),
};
let tool_call_id = tc_id.clone();
let interrupt_message = message.clone();
let interrupt_id = uuid::Uuid::new_v4().to_string();
let open_interrupt = OpenInterrupt {
interrupt_id: interrupt_id.clone(),
tool_call_id: tool_call_id.clone(),
reason: "tool_call".into(),
message: Some(interrupt_message.clone()),
created_at: crate::session::chrono_lite_now(),
};
if let Some(session_dir) = agui_session_dir(&drv_workspace, &drv_thread) {
let _ = std::fs::create_dir_all(&session_dir);
save_open_interrupts(&session_dir, std::slice::from_ref(&open_interrupt));
if let Ok(state_val) = serde_json::to_value(runtime.transcript()) {
let _ = sse_tx.send(ag::Event::StateSnapshot(ag::StateSnapshot {
snapshot: state_val,
base: ag::BaseEvent::default(),
}));
}
let messages_json: Vec<serde_json::Value> = runtime
.transcript()
.iter()
.filter_map(|m| serde_json::to_value(m).ok())
.collect();
let _ = sse_tx.send(ag::Event::MessagesSnapshot(ag::MessagesSnapshot {
messages: messages_json,
base: ag::BaseEvent::default(),
}));
}
let _ = sse_tx.send(ag::Event::RunFinished(ag::RunFinished {
thread_id: drv_thread,
run_id: drv_run,
outcome: Some(ag::RunFinishedOutcome::Interrupt {
interrupts: vec![ag::Interrupt {
id: interrupt_id,
reason: "tool_call".into(),
message: Some(interrupt_message),
tool_call_id: Some(tool_call_id),
response_schema: Some(response_schema),
expires_at: None,
metadata: Some(metadata),
}],
}),
result: None,
base: ag::BaseEvent::default(),
}));
} else {
let (run_outcome, result_msg) = match &outcome {
Ok(o) => (
ag::RunFinishedOutcome::Success,
o.final_text.clone().map(serde_json::Value::String),
),
Err(e) => (
ag::RunFinishedOutcome::Error {
message: e.to_string(),
code: None,
},
None,
),
};
let _ = sse_tx.send(ag::Event::RunFinished(ag::RunFinished {
thread_id: drv_thread,
run_id: drv_run,
outcome: Some(run_outcome),
result: result_msg,
base: ag::BaseEvent::default(),
}));
}
});
tokio::spawn(async move {
if let Err(e) = driver_handle.await {
tracing::error!("agui: driver task panicked: {e}");
}
});
let stream = tokio_stream::wrappers::UnboundedReceiverStream::new(sse_rx).map(|ev| {
let data = serde_json::to_string(&ev).unwrap_or_else(|_| "{}".into());
Ok::<_, Infallible>(Event::default().data(data))
});
Ok(Sse::new(stream))
}
fn sanitize_thread_id_for_session(thread: &str) -> String {
let mut out: String = thread
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' {
c
} else {
'-'
}
})
.collect();
while out.starts_with('.') {
out.replace_range(..1, "-");
}
while out.contains("..") {
out = out.replace("..", "-.");
}
if out.is_empty() {
out.push_str("default");
}
out
}
fn agui_session_dir(workspace: &std::path::Path, thread_id: &str) -> Option<std::path::PathBuf> {
let session_id = sanitize_thread_id_for_session(thread_id);
crate::user_sessions_dir(workspace)
.ok()
.map(|d| d.join(format!("agui-{session_id}")))
}
#[derive(Debug, Clone)]
enum AguiInterruptDetail {
Test {
tool_call_id: String,
tool_name: String,
},
Client {
tool_call_id: String,
tool_name: String,
parameters: serde_json::Value,
args: serde_json::Value,
},
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
struct OpenInterrupt {
pub interrupt_id: String,
pub tool_call_id: String,
pub reason: String,
pub message: Option<String>,
pub created_at: String,
}
fn load_open_interrupts(session_dir: &std::path::Path) -> Vec<OpenInterrupt> {
let path = session_dir.join(".interrupts.json");
let Ok(bytes) = std::fs::read(&path) else {
return Vec::new();
};
serde_json::from_slice(&bytes).unwrap_or_default()
}
fn save_open_interrupts(session_dir: &std::path::Path, interrupts: &[OpenInterrupt]) {
if let Ok(json) = serde_json::to_string_pretty(interrupts) {
let path = session_dir.join(".interrupts.json");
let _ = std::fs::create_dir_all(session_dir);
crate::atomic::atomic_write(&path, json.as_bytes()).ok();
}
}
fn clear_open_interrupts(session_dir: &std::path::Path) {
let path = session_dir.join(".interrupts.json");
let _ = std::fs::remove_file(&path);
}
struct TestInterruptHook {
interrupt_before: Vec<String>,
interrupted_tool_call_id: std::sync::Mutex<Option<String>>,
interrupted_tool_name: std::sync::Mutex<Option<String>>,
interrupted_arguments: std::sync::Mutex<Option<serde_json::Value>>,
}
pub(crate) const CLIENT_TOOL_DENY_PREFIX: &str = "[frontend tool]";
struct ClientToolStub {
name: String,
description: String,
parameters: serde_json::Value,
}
#[async_trait::async_trait]
impl crate::tools::Tool for ClientToolStub {
fn spec(&self) -> crate::llm::chat::ToolSpec {
crate::llm::chat::ToolSpec {
name: self.name.clone(),
description: self.description.clone(),
parameters: self.parameters.clone(),
}
}
async fn execute(&self, _arguments: serde_json::Value) -> crate::error::Result<String> {
Ok(format!(
"{} {} not executed server-side",
CLIENT_TOOL_DENY_PREFIX, self.name
))
}
}
struct ClientToolHook {
names: std::collections::HashSet<String>,
denied: std::sync::Mutex<Option<(String, String)>>,
}
#[async_trait::async_trait]
impl crate::tools::PermissionHook for ClientToolHook {
async fn check(
&self,
name: &str,
args: &serde_json::Value,
) -> crate::agent::PermissionDecision {
if !self.names.contains(name) {
return crate::agent::PermissionDecision::Allow;
}
{
let mut slot = self.denied.lock().unwrap_or_else(|e| e.into_inner());
if slot.is_none() {
*slot = Some((name.to_string(), args.to_string()));
}
}
crate::agent::PermissionDecision::Deny(format!(
"{} {} awaiting client execution",
CLIENT_TOOL_DENY_PREFIX, name
))
}
}
#[async_trait::async_trait]
#[async_trait]
impl crate::tools::PermissionHook for TestInterruptHook {
async fn check(
&self,
name: &str,
args: &serde_json::Value,
) -> crate::agent::PermissionDecision {
if self
.interrupted_tool_call_id
.lock()
.unwrap_or_else(|e| e.into_inner())
.is_some()
{
return crate::agent::PermissionDecision::Allow;
}
if self.interrupt_before.iter().any(|n| n == name) {
*self
.interrupted_tool_name
.lock()
.unwrap_or_else(|e| e.into_inner()) = Some(name.to_string());
*self
.interrupted_arguments
.lock()
.unwrap_or_else(|e| e.into_inner()) = Some(args.clone());
return crate::agent::PermissionDecision::Deny(
"test interrupt trigger — tool blocked by interrupt_before".into(),
);
}
crate::agent::PermissionDecision::Allow
}
}
pub(super) async fn metrics_handler(State(state): State<Arc<AppState>>) -> String {
let metrics = &state.metrics;
let requests_total = metrics.requests_total.load(Ordering::Relaxed);
let requests_active = metrics.requests_active.load(Ordering::Relaxed);
let agent_runs_total = metrics.agent_runs_total.load(Ordering::Relaxed);
let agent_runs_success = metrics.agent_runs_success.load(Ordering::Relaxed);
let agent_runs_failed = metrics.agent_runs_failed.load(Ordering::Relaxed);
let tokens_prompt_total = metrics.tokens_prompt_total.load(Ordering::Relaxed);
let tokens_completion_total = metrics.tokens_completion_total.load(Ordering::Relaxed);
let agent_steps_total = metrics.agent_steps_total.load(Ordering::Relaxed);
let sessions_active = metrics.sessions_active.load(Ordering::Relaxed);
let rate_limits_rejected = metrics.rate_limits_rejected.load(Ordering::Relaxed);
format!(
"# HELP recursive_requests_total Total HTTP requests\n\
# TYPE recursive_requests_total counter\n\
recursive_requests_total {requests_total}\n\
# HELP recursive_requests_active Currently active HTTP requests\n\
# TYPE recursive_requests_active gauge\n\
recursive_requests_active {requests_active}\n\
# HELP recursive_agent_runs_total Total agent runs\n\
# TYPE recursive_agent_runs_total counter\n\
recursive_agent_runs_total {agent_runs_total}\n\
# HELP recursive_agent_runs_success Successful agent runs\n\
# TYPE recursive_agent_runs_success counter\n\
recursive_agent_runs_success {agent_runs_success}\n\
# HELP recursive_agent_runs_failed Failed agent runs\n\
# TYPE recursive_agent_runs_failed counter\n\
recursive_agent_runs_failed {agent_runs_failed}\n\
# HELP recursive_tokens_prompt_total Total prompt tokens consumed\n\
# TYPE recursive_tokens_prompt_total counter\n\
recursive_tokens_prompt_total {tokens_prompt_total}\n\
# HELP recursive_tokens_completion_total Total completion tokens generated\n\
# TYPE recursive_tokens_completion_total counter\n\
recursive_tokens_completion_total {tokens_completion_total}\n\
# HELP recursive_agent_steps_total Total agent steps executed\n\
# TYPE recursive_agent_steps_total counter\n\
recursive_agent_steps_total {agent_steps_total}\n\
# HELP recursive_sessions_active Currently active sessions\n\
# TYPE recursive_sessions_active gauge\n\
recursive_sessions_active {sessions_active}\n\
# HELP recursive_rate_limits_rejected_total Total requests rejected by rate limiting\n\
# TYPE recursive_rate_limits_rejected_total counter\n\
recursive_rate_limits_rejected_total {rate_limits_rejected}\n"
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::event::AgentEvent;
use crate::http::SseEvent;
#[test]
fn agui_events_for_assistant_text_emits_start_content_end() {
use agui_protocol as ag;
let ev = AgentEvent::AssistantText {
text: "hi".into(),
step: 0,
};
let out = agui_events_for(&ev);
assert_eq!(out.len(), 3, "got {out:?}");
assert!(matches!(out[0], ag::Event::TextMessageStart(_)));
assert!(matches!(out[1], ag::Event::TextMessageContent(_)));
assert!(matches!(out[2], ag::Event::TextMessageEnd(_)));
}
#[test]
fn tool_progress_emitted_after_tool_result() {
use std::collections::HashMap;
use std::time::Instant;
let mut tool_start_times: HashMap<String, Instant> = HashMap::new();
let mut emitted: Vec<SseEvent> = Vec::new();
let call_event = AgentEvent::ToolCall {
name: "Bash".to_string(),
id: "tc-1".to_string(),
arguments: "{}".to_string(),
step: 0,
};
if let AgentEvent::ToolCall { id, .. } = &call_event {
tool_start_times.insert(id.clone(), Instant::now());
}
if let Some(ev) = map_agent_event(&call_event) {
emitted.push(ev);
}
let result_event = AgentEvent::ToolResult {
id: "tc-1".to_string(),
name: "Bash".to_string(),
output: "ok".to_string(),
step: 0,
is_error: false,
};
if let Some(ev) = map_agent_event(&result_event) {
emitted.push(ev);
}
if let AgentEvent::ToolResult { id, name, .. } = &result_event {
let elapsed_ms = tool_start_times
.remove(id)
.map(|start| start.elapsed().as_millis() as u64)
.unwrap_or(0);
emitted.push(SseEvent::ToolProgress {
tool_use_id: id.clone(),
tool_name: name.clone(),
elapsed_ms,
});
}
assert_eq!(emitted.len(), 3, "expected 3 events");
assert!(matches!(emitted[0], SseEvent::ToolCall { .. }));
assert!(matches!(emitted[1], SseEvent::ToolResult { .. }));
let SseEvent::ToolProgress {
tool_use_id,
tool_name,
elapsed_ms,
} = &emitted[2]
else {
panic!("third event should be ToolProgress");
};
assert_eq!(tool_use_id, "tc-1");
assert_eq!(tool_name, "Bash");
let _ = elapsed_ms; }
#[test]
fn tool_progress_elapsed_is_zero_for_unmatched_result() {
use std::collections::HashMap;
use std::time::Instant;
let mut tool_start_times: HashMap<String, Instant> = HashMap::new();
let result_event = AgentEvent::ToolResult {
id: "tc-orphan".to_string(),
name: "Read".to_string(),
output: "data".to_string(),
step: 0,
is_error: false,
};
let elapsed_ms = if let AgentEvent::ToolResult { id, .. } = &result_event {
tool_start_times
.remove(id)
.map(|start| start.elapsed().as_millis() as u64)
.unwrap_or(0)
} else {
unreachable!()
};
assert_eq!(elapsed_ms, 0);
}
#[tokio::test]
async fn agui_run_respects_run_semaphore() {
use crate::llm::MockProvider;
use crate::tools::ToolRegistry;
use std::sync::Arc;
use tokio::sync::Semaphore;
std::env::set_var("RECURSIVE_API_KEY", "test-key");
std::env::set_var("RECURSIVE_MODEL", "test-model");
let config = crate::config::Config::from_env().unwrap();
let sem = Arc::new(Semaphore::new(0));
let state = Arc::new(crate::http::AppState {
tools: vec![],
tool_registry: ToolRegistry::default(),
config,
provider: Arc::new(MockProvider::new(vec![])),
sessions: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
event_channels: Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())),
metrics: Arc::new(crate::http::Metrics::default()),
slash_commands: Arc::new(vec![]),
session_ttl_secs: 3600,
run_semaphore: sem,
rate_limiter: crate::http::RateLimiter::new(10, 1.0),
skills: vec![],
});
let body = serde_json::json!({
"threadId": "t1",
"runId": "r1",
"messages": [{"id": "m1", "role": "user", "content": "hi"}],
});
let (status, _err) = agui_run(State(state), Json(body))
.await
.expect_err("expected SERVICE_UNAVAILABLE");
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
}
#[tokio::test]
async fn clear_goal_returns_409_when_runtime_busy() {
use crate::llm::MockProvider;
use crate::tools::ToolRegistry;
use axum::response::IntoResponse;
use std::sync::Arc;
use tokio::sync::Semaphore;
std::env::set_var("RECURSIVE_API_KEY", "test-key");
std::env::set_var("RECURSIVE_MODEL", "test-model");
let config = crate::config::Config::from_env().unwrap();
let provider = Arc::new(MockProvider::new(vec![]));
let session_id = "test-busy-session".to_string();
let runtime = AgentRuntimeBuilder::new()
.llm(provider.clone())
.tools(ToolRegistry::default())
.build()
.expect("runtime build");
let runtime_arc = Arc::new(tokio::sync::Mutex::new(runtime));
let session = SessionState {
id: session_id.clone(),
created_at: "2026-01-01T00:00:00Z".to_string(),
title: None,
runtime: runtime_arc.clone(),
plan_approval_gate: Default::default(),
interrupt_token: Arc::new(tokio::sync::Mutex::new(None)),
non_system_message_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
last_active_ms: Arc::new(AtomicU64::new(0)),
prompt_tokens: Arc::new(AtomicU64::new(0)),
completion_tokens: Arc::new(AtomicU64::new(0)),
};
let sessions: HashMap<String, SessionState> = [(session_id.clone(), session)].into();
let state = Arc::new(AppState {
tools: vec![],
tool_registry: ToolRegistry::default(),
config,
provider,
sessions: Arc::new(tokio::sync::RwLock::new(sessions)),
event_channels: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
metrics: Arc::new(crate::http::Metrics::default()),
slash_commands: Arc::new(vec![]),
session_ttl_secs: 3600,
run_semaphore: Arc::new(Semaphore::new(8)),
rate_limiter: crate::http::RateLimiter::new(10, 1.0),
skills: vec![],
});
let guard = runtime_arc.lock().await;
let resp = session_clear_goal(State(state.clone()), Path(session_id.clone()))
.await
.into_response();
let status = resp.status();
assert_eq!(status, StatusCode::CONFLICT, "expected 409 Conflict");
let retry_after = resp
.headers()
.get(axum::http::header::RETRY_AFTER)
.expect("Retry-After header missing")
.to_str()
.unwrap();
assert_eq!(retry_after, "5", "expected Retry-After: 5");
drop(guard);
let resp = session_clear_goal(State(state), Path(session_id))
.await
.into_response();
assert_eq!(resp.status(), StatusCode::OK, "expected 200 after unlock");
}
fn test_app_state_with_session(
session_id: &str,
) -> (
Arc<AppState>,
Arc<tokio::sync::Mutex<crate::runtime::AgentRuntime>>,
) {
use crate::llm::MockProvider;
use crate::tools::ToolRegistry;
use tokio::sync::Semaphore;
std::env::set_var("RECURSIVE_API_KEY", "test-key");
std::env::set_var("RECURSIVE_MODEL", "test-model");
let config = crate::config::Config::from_env().unwrap();
let provider = Arc::new(MockProvider::new(vec![]));
let runtime = AgentRuntimeBuilder::new()
.llm(provider.clone())
.tools(ToolRegistry::default())
.build()
.expect("runtime build");
let runtime_arc = Arc::new(tokio::sync::Mutex::new(runtime));
let session = SessionState {
id: session_id.to_string(),
created_at: "2026-01-01T00:00:00Z".to_string(),
title: Some("old".into()),
runtime: runtime_arc.clone(),
plan_approval_gate: Default::default(),
interrupt_token: Arc::new(tokio::sync::Mutex::new(None)),
non_system_message_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
last_active_ms: Arc::new(AtomicU64::new(0)),
prompt_tokens: Arc::new(AtomicU64::new(0)),
completion_tokens: Arc::new(AtomicU64::new(0)),
};
let sessions: HashMap<String, SessionState> = [(session_id.to_string(), session)].into();
let state = Arc::new(AppState {
tools: vec![],
tool_registry: ToolRegistry::default(),
config,
provider,
sessions: Arc::new(tokio::sync::RwLock::new(sessions)),
event_channels: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
metrics: Arc::new(crate::http::Metrics::default()),
slash_commands: Arc::new(vec![]),
session_ttl_secs: 3600,
run_semaphore: Arc::new(Semaphore::new(8)),
rate_limiter: crate::http::RateLimiter::new(10, 1.0),
skills: vec![],
});
(state, runtime_arc)
}
#[tokio::test]
async fn get_session_status_idle_vs_plan_pending() {
let sid = "test-plan-status";
let (state, _) = test_app_state_with_session(sid);
let idle = match get_session(State(state.clone()), Path(sid.to_string())).await {
Ok(Json(v)) => v,
Err(_) => panic!("get_session idle"),
};
assert_eq!(idle.status, "idle");
assert!(idle.pending_plan.is_none());
{
let sessions = state.sessions.read().await;
let session = sessions.get(sid).unwrap();
*session.plan_approval_gate.pending_plan.write().unwrap() = Some("do the thing".into());
}
let pending = match get_session(State(state), Path(sid.to_string())).await {
Ok(Json(v)) => v,
Err(_) => panic!("get_session pending"),
};
assert_eq!(pending.status, "plan_pending_approval");
assert_eq!(pending.pending_plan.as_deref(), Some("do the thing"));
}
#[tokio::test]
async fn get_session_busy_runtime_returns_empty_messages() {
let sid = "test-busy-get";
let (state, runtime_arc) = test_app_state_with_session(sid);
let _guard = runtime_arc.lock().await;
let detail = match get_session(State(state), Path(sid.to_string())).await {
Ok(Json(v)) => v,
Err(_) => panic!("busy get_session must still 200"),
};
assert!(
detail.messages.is_empty(),
"busy runtime must fall back to empty messages"
);
assert_eq!(detail.status, "idle");
}
#[tokio::test]
async fn patch_session_empty_title_clears() {
let sid = "test-patch-title";
let (state, _) = test_app_state_with_session(sid);
let cleared = match patch_session(
State(state.clone()),
Path(sid.to_string()),
Json(PatchSessionRequest {
title: Some("".into()),
}),
)
.await
{
Ok(Json(v)) => v,
Err(_) => panic!("patch empty title"),
};
assert!(cleared.title.is_none(), "empty title must clear to None");
let set = match patch_session(
State(state.clone()),
Path(sid.to_string()),
Json(PatchSessionRequest {
title: Some("new".into()),
}),
)
.await
{
Ok(Json(v)) => v,
Err(_) => panic!("patch new title"),
};
assert_eq!(set.title.as_deref(), Some("new"));
let keep = match patch_session(
State(state),
Path(sid.to_string()),
Json(PatchSessionRequest { title: None }),
)
.await
{
Ok(Json(v)) => v,
Err(_) => panic!("patch omit title"),
};
assert_eq!(keep.title.as_deref(), Some("new"));
}
#[tokio::test]
async fn metrics_handler_includes_new_fields() {
use crate::http::Metrics;
use crate::tools::ToolRegistry;
use std::sync::atomic::AtomicU64;
std::env::set_var("RECURSIVE_API_KEY", "test-key");
std::env::set_var("RECURSIVE_MODEL", "test-model");
let config = crate::config::Config::from_env().unwrap();
let metrics = Metrics {
sessions_active: AtomicU64::new(3),
rate_limits_rejected: AtomicU64::new(42),
..Metrics::default()
};
let state = Arc::new(AppState {
metrics: Arc::new(metrics),
tools: vec![],
tool_registry: ToolRegistry::default(),
config,
provider: Arc::new(crate::llm::MockProvider::new(vec![])),
sessions: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
event_channels: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
slash_commands: Arc::new(vec![]),
session_ttl_secs: 3600,
run_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
rate_limiter: crate::http::RateLimiter::new(10, 1.0),
skills: vec![],
});
let output = metrics_handler(State(state)).await;
assert!(
output.contains("recursive_sessions_active 3"),
"output should contain sessions_active: {output}"
);
assert!(
output.contains("recursive_rate_limits_rejected_total 42"),
"output should contain rate_limits_rejected_total: {output}"
);
}
#[tokio::test]
async fn sessions_active_tracks_session_lifecycle() {
use crate::tools::ToolRegistry;
use tower::ServiceExt;
std::env::set_var("RECURSIVE_API_KEY", "test-key");
std::env::set_var("RECURSIVE_MODEL", "test-model");
std::env::set_var("RECURSIVE_HTTP_AUTH_INSECURE_OK", "1");
let config = crate::config::Config::from_env().unwrap();
let provider = Arc::new(crate::llm::MockProvider::new(vec![]));
let metrics = Arc::new(crate::http::Metrics::default());
let state = AppState {
tools: vec![],
tool_registry: ToolRegistry::default(),
config,
provider,
sessions: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
event_channels: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
metrics: metrics.clone(),
slash_commands: Arc::new(vec![]),
session_ttl_secs: 3600,
run_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
rate_limiter: crate::http::RateLimiter::new(100, 1.0),
skills: vec![],
};
let auth = crate::http::auth::AuthConfig::default();
let limiter = state.rate_limiter.clone();
let app = crate::http::build_router_with_auth_and_rate_limit(state, auth, limiter);
assert_eq!(
metrics.sessions_active.load(Ordering::Relaxed),
0,
"sessions_active should start at 0"
);
let resp = app
.clone()
.oneshot(
axum::extract::Request::builder()
.method("POST")
.uri("/sessions")
.header("content-type", "application/json")
.body(axum::body::Body::from(r#"{"session_name":"test"}"#))
.unwrap(),
)
.await
.unwrap();
assert_eq!(
resp.status(),
StatusCode::CREATED,
"expected 201 Created from POST /sessions"
);
let body = axum::body::to_bytes(resp.into_body(), 1024).await.unwrap();
let created: CreateSessionResponse =
serde_json::from_slice(&body).expect("valid CreateSessionResponse");
let session_id = created.id;
assert_eq!(
metrics.sessions_active.load(Ordering::Relaxed),
1,
"sessions_active should increment to 1 after create"
);
let resp = app
.clone()
.oneshot(
axum::extract::Request::builder()
.method("DELETE")
.uri(format!("/sessions/{session_id}"))
.body(axum::body::Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(
resp.status(),
StatusCode::NO_CONTENT,
"expected 204 No Content from DELETE /sessions/:id"
);
assert_eq!(
metrics.sessions_active.load(Ordering::Relaxed),
0,
"sessions_active should decrement to 0 after delete"
);
}
#[test]
fn openapi_session_detail_has_complete_schema() {
let spec = super::build_openapi_spec();
let props = &spec["components"]["schemas"]["SessionDetailResponse"]["properties"];
assert!(props.get("id").is_some(), "id missing");
assert!(props.get("created_at").is_some(), "created_at missing");
assert!(props.get("messages").is_some(), "messages missing");
assert!(
props.get("prompt_tokens").is_some(),
"prompt_tokens missing"
);
assert!(
props.get("completion_tokens").is_some(),
"completion_tokens missing"
);
assert!(props.get("status").is_some(), "status missing");
assert!(props.get("todos").is_some(), "todos missing");
assert!(props.get("goal").is_some(), "goal missing");
assert!(props.get("title").is_some(), "title missing");
assert!(props.get("pending_plan").is_some(), "pending_plan missing");
assert!(props.get("first_prompt").is_some(), "first_prompt missing");
assert!(props.get("last_prompt").is_some(), "last_prompt missing");
let obj = props.as_object().expect("properties is an object");
assert!(
obj.len() >= 10,
"SessionDetailResponse should have ≥10 properties, got {}",
obj.len()
);
}
#[test]
fn openapi_session_info_has_message_count_and_title() {
let spec = super::build_openapi_spec();
let props = &spec["components"]["schemas"]["SessionInfo"]["properties"];
assert!(props.get("id").is_some(), "id missing");
assert!(props.get("created_at").is_some(), "created_at missing");
assert!(
props.get("message_count").is_some(),
"message_count missing"
);
assert!(props.get("title").is_some(), "title missing");
}
#[test]
fn openapi_metrics_path_documents_new_metrics() {
let spec = super::build_openapi_spec();
let description = spec["paths"]["/metrics"]["get"]["description"]
.as_str()
.expect("metrics path description is a string");
assert!(
description.contains("recursive_sessions_active"),
"metrics description should mention recursive_sessions_active: {description}"
);
assert!(
description.contains("recursive_rate_limits_rejected_total"),
"metrics description should mention recursive_rate_limits_rejected_total: {description}"
);
}
#[test]
fn list_sessions_sort_is_chronological() {
let mut infos = [
SessionInfo {
id: "c".into(),
created_at: "2026-01-03T00:00:00Z".into(),
message_count: 0,
title: None,
},
SessionInfo {
id: "a".into(),
created_at: "2026-01-01T00:00:00Z".into(),
message_count: 0,
title: None,
},
SessionInfo {
id: "b".into(),
created_at: "2026-01-02T00:00:00Z".into(),
message_count: 0,
title: None,
},
];
infos.sort_by(|a, b| a.created_at.cmp(&b.created_at).then(a.id.cmp(&b.id)));
assert_eq!(infos[0].id, "a");
assert_eq!(infos[1].id, "b");
assert_eq!(infos[2].id, "c");
}
#[test]
fn list_sessions_same_second_tiebreak_by_id() {
let mut infos = [
SessionInfo {
id: "z".into(),
created_at: "2026-01-01T00:00:00Z".into(),
message_count: 0,
title: None,
},
SessionInfo {
id: "a".into(),
created_at: "2026-01-01T00:00:00Z".into(),
message_count: 0,
title: None,
},
];
infos.sort_by(|a, b| a.created_at.cmp(&b.created_at).then(a.id.cmp(&b.id)));
assert_eq!(infos[0].id, "a");
assert_eq!(infos[1].id, "z");
}
#[test]
fn skill_reminder_produced_and_not_inlined() {
let skills = vec![crate::skills::Skill {
name: "rust-patch-discipline".to_string(),
description: "V4A patch format rules".to_string(),
path: std::path::PathBuf::from("/tmp/skills/rust-patch-discipline/SKILL.md"),
mode: crate::skills::SkillMode::Manual,
triggers: vec![],
hint: String::new(),
depends_on: vec![],
refs: vec![],
params: vec![],
scripts: vec![],
sections: vec![],
globs: None,
}];
let idx = crate::skills::skill_index(&skills);
assert!(!idx.is_empty(), "skill_index should not be empty");
assert!(
idx.contains("Available skills"),
"skill_index should contain header"
);
assert!(
idx.contains("rust-patch-discipline"),
"skill_index should list skill names"
);
let reminder = crate::skills::skill_reminder(&skills);
assert!(reminder.contains("<system-reminder>"), "reminder wrapper");
assert!(
reminder.contains("Available skills"),
"reminder should contain header"
);
assert!(
reminder.contains("rust-patch-discipline"),
"reminder should contain skill name"
);
let base_prompt = "You are a helpful agent.";
let assembled =
crate::assemble_system_prompt(base_prompt, std::path::Path::new(""), &skills, false)
.into_full();
assert!(
!assembled.contains("Available skills"),
"system prompt must not inline skill catalog: {assembled}"
);
}
#[test]
fn system_prompt_unchanged_when_no_skills() {
let skills: Vec<crate::skills::Skill> = vec![];
let idx = crate::skills::skill_index(&skills);
assert!(idx.is_empty(), "empty skills should produce empty index");
let base_prompt = "You are a helpful agent.";
let mut sp = base_prompt.to_string();
if !idx.is_empty() {
sp.push('\n');
sp.push_str(&idx);
}
assert_eq!(sp, base_prompt);
}
#[test]
fn parse_permission_mode_all_variants() {
assert_eq!(parse_permission_mode("auto", false), PermissionMode::Auto);
assert_eq!(parse_permission_mode("AUTO", true), PermissionMode::Auto);
assert_eq!(
parse_permission_mode("strict", false),
PermissionMode::Strict
);
assert_eq!(
parse_permission_mode("bypass", true),
PermissionMode::BypassPermissions
);
assert_eq!(
parse_permission_mode("bypass_permissions", true),
PermissionMode::BypassPermissions
);
assert_eq!(
parse_permission_mode("bypass", false),
PermissionMode::Default
);
assert_eq!(
parse_permission_mode("bypass_permissions", false),
PermissionMode::Default
);
assert_eq!(
parse_permission_mode("default", true),
PermissionMode::Default
);
assert_eq!(
parse_permission_mode("unknown", true),
PermissionMode::Default
);
}
#[test]
fn map_agent_event_tool_result_success_flag() {
let ok = AgentEvent::ToolResult {
id: "tc-1".into(),
name: "Bash".into(),
output: "ok".into(),
step: 0,
is_error: false,
};
let err = AgentEvent::ToolResult {
id: "tc-2".into(),
name: "Bash".into(),
output: "fail".into(),
step: 0,
is_error: true,
};
match map_agent_event(&ok) {
Some(SseEvent::ToolResult { success, .. }) => assert!(success),
other => panic!("expected ToolResult success=true, got {other:?}"),
}
match map_agent_event(&err) {
Some(SseEvent::ToolResult { success, .. }) => assert!(!success),
other => panic!("expected ToolResult success=false, got {other:?}"),
}
}
#[test]
fn map_agent_event_suppresses_non_sse_variants() {
assert!(map_agent_event(&AgentEvent::Latency { step: 0, llm_ms: 1 }).is_none());
assert!(map_agent_event(&AgentEvent::Usage {
input_tokens: 1,
output_tokens: 1,
cache_hit_tokens: 0,
cache_miss_tokens: 0,
step: 0,
})
.is_none());
assert!(map_agent_event(&AgentEvent::AssistantText {
text: "hi".into(),
step: 0
})
.is_none());
}
#[test]
fn map_agent_event_forwards_goal_loop_events() {
match map_agent_event(&AgentEvent::GoalContinuing {
reason: "still working".into(),
turns: 3,
}) {
Some(SseEvent::GoalContinuing { reason, turns }) => {
assert_eq!(reason, "still working");
assert_eq!(turns, 3);
}
other => panic!("expected GoalContinuing, got {other:?}"),
}
match map_agent_event(&AgentEvent::GoalAchieved {
condition: "tests pass".into(),
turns: 5,
}) {
Some(SseEvent::GoalAchieved { condition, turns }) => {
assert_eq!(condition, "tests pass");
assert_eq!(turns, 5);
}
other => panic!("expected GoalAchieved, got {other:?}"),
}
}
#[test]
fn map_agent_event_forwards_core_sse_arms() {
match map_agent_event(&AgentEvent::PartialToken {
text: "tok".into(),
step: 2,
}) {
Some(SseEvent::PartialMessage { text, step }) => {
assert_eq!(text, "tok");
assert_eq!(step, 2);
}
other => panic!("expected PartialMessage, got {other:?}"),
}
match map_agent_event(&AgentEvent::ToolCall {
name: "Bash".into(),
id: "tc-1".into(),
arguments: "{}".into(),
step: 1,
}) {
Some(SseEvent::ToolCall { name, step }) => {
assert_eq!(name, "Bash");
assert_eq!(step, 1);
}
other => panic!("expected ToolCall, got {other:?}"),
}
match map_agent_event(&AgentEvent::TurnFinished {
reason: "no_more_tool_calls".into(),
steps: 4,
}) {
Some(SseEvent::Done {
finish_reason,
total_steps,
}) => {
assert_eq!(finish_reason, "no_more_tool_calls");
assert_eq!(total_steps, 4);
}
other => panic!("expected Done, got {other:?}"),
}
match map_agent_event(&AgentEvent::PlanProposed {
plan_text: "do X".into(),
tool_calls: vec![],
}) {
Some(SseEvent::PlanProposed { plan }) => assert_eq!(plan, "do X"),
other => panic!("expected PlanProposed, got {other:?}"),
}
}
#[test]
fn sse_message_from_canonical_filters_system_tool_and_empty() {
assert!(
sse_message_from_canonical(&crate::message::Message::system("seed")).is_none(),
"system messages must be filtered"
);
let mut tool = crate::message::Message::user("unused");
tool.role = crate::message::Role::Tool;
tool.tool_call_id = Some("tc-1".into());
assert!(
sse_message_from_canonical(&tool).is_none(),
"tool messages must be filtered"
);
assert!(
sse_message_from_canonical(&crate::message::Message::assistant("")).is_none(),
"empty assistant with no tool_calls must be None"
);
match sse_message_from_canonical(&crate::message::Message::assistant("hello")) {
Some(SseEvent::Message { role, content }) => {
assert_eq!(role, "assistant");
assert!(!content.is_empty());
}
other => panic!("expected Message event, got {other:?}"),
}
match sse_message_from_canonical(&crate::message::Message::user("hi")) {
Some(SseEvent::Message { role, .. }) => assert_eq!(role, "user"),
other => panic!("expected user Message, got {other:?}"),
}
}
#[test]
fn sse_message_from_canonical_emits_tool_use_without_text() {
let msg = crate::message::Message::assistant_with_tool_calls(
"",
vec![crate::llm::ToolCall {
id: "tc-9".into(),
name: "Read".into(),
arguments: serde_json::json!({"path":"a.rs"}),
}],
);
match sse_message_from_canonical(&msg) {
Some(SseEvent::Message { role, content }) => {
assert_eq!(role, "assistant");
assert!(
content.iter().any(|b| matches!(
b,
SseContentBlock::ToolUse { id, name, .. }
if id == "tc-9" && name == "Read"
)),
"expected ToolUse block, got {content:?}"
);
}
other => panic!("expected Message with ToolUse, got {other:?}"),
}
}
#[test]
fn agui_events_for_partial_token_then_tool_call_closes_stream() {
use agui_protocol as ag;
let mut conv = AguiConverter::new();
let t1 = conv.convert(&AgentEvent::PartialToken {
text: "a".into(),
step: 0,
});
assert_eq!(t1.len(), 2, "Start+Content expected: {t1:?}");
assert!(matches!(t1[0], ag::Event::TextMessageStart(_)));
assert!(matches!(t1[1], ag::Event::TextMessageContent(_)));
let t2 = conv.convert(&AgentEvent::PartialToken {
text: "b".into(),
step: 0,
});
assert_eq!(t2.len(), 1, "Content-only expected: {t2:?}");
assert!(matches!(t2[0], ag::Event::TextMessageContent(_)));
let t3 = conv.convert(&AgentEvent::ToolCall {
id: "tc-1".into(),
name: "Bash".into(),
arguments: "{}".into(),
step: 0,
});
assert!(
t3.iter().any(|e| matches!(e, ag::Event::TextMessageEnd(_))),
"ToolCall must close open text stream: {t3:?}"
);
assert!(
t3.iter().any(|e| matches!(e, ag::Event::ToolCallStart(_))),
"ToolCall must emit ToolCallStart: {t3:?}"
);
assert!(
t3.iter().any(|e| matches!(e, ag::Event::ToolCallArgs(_))),
"ToolCall must emit ToolCallArgs: {t3:?}"
);
assert!(
t3.iter().any(|e| matches!(e, ag::Event::ToolCallEnd(_))),
"ToolCall must emit ToolCallEnd: {t3:?}"
);
}
#[test]
fn format_timestamp_unix_epoch_is_1970() {
let ts = format_timestamp(SystemTime::UNIX_EPOCH);
assert_eq!(
&ts[..10],
"1970-01-01",
"epoch date must be 1970-01-01; got {ts}"
);
assert_eq!(
&ts[11..19],
"00:00:00",
"epoch time must be 00:00:00; got {ts}"
);
assert!(ts.ends_with('Z'), "must end with Z; got {ts}");
}
#[test]
fn format_timestamp_known_date() {
let secs = 1_783_296_000u64;
let t = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(secs);
let ts = format_timestamp(t);
assert_eq!(&ts[..10], "2026-07-06", "expected 2026-07-06; got {ts}");
assert_eq!(&ts[11..19], "00:00:00", "time part must be zero; got {ts}");
}
#[test]
fn format_timestamp_time_parts_in_range() {
let secs: u64 = 1_783_296_000 + 13 * 3600 + 45 * 60 + 30;
let t = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(secs);
let ts = format_timestamp(t);
assert_eq!(&ts[11..13], "13", "hours mismatch; got {ts}");
assert_eq!(&ts[14..16], "45", "minutes mismatch; got {ts}");
assert_eq!(&ts[17..19], "30", "seconds mismatch; got {ts}");
}
#[test]
fn sanitize_thread_id_valid_passthrough() {
assert_eq!(sanitize_thread_id_for_session("abc-123"), "abc-123");
assert_eq!(sanitize_thread_id_for_session("foo_bar.baz"), "foo_bar.baz");
}
#[test]
fn sanitize_thread_id_replaces_special_chars() {
let out = sanitize_thread_id_for_session("a/b:c");
assert!(!out.contains('/'), "slash must be replaced");
assert!(!out.contains(':'), "colon must be replaced");
}
#[test]
fn sanitize_thread_id_leading_dot_replaced() {
let out = sanitize_thread_id_for_session(".hidden");
assert!(
!out.starts_with('.'),
"leading dot must be replaced; got {out}"
);
}
#[test]
fn sanitize_thread_id_double_dot_collapsed() {
let out = sanitize_thread_id_for_session("a..b");
assert!(
!out.contains(".."),
"double dot must be collapsed; got {out}"
);
}
#[test]
fn sanitize_thread_id_empty_becomes_default() {
assert_eq!(sanitize_thread_id_for_session(""), "default");
}
}