use a8e_core::agents::extension::{Envs, PLATFORM_EXTENSIONS};
use a8e_core::agents::{A8ePlatform, Agent, AgentConfig, ExtensionConfig, SessionConfig};
use a8e_core::builtin_extension::register_builtin_extensions;
use a8e_core::config::base::CONFIG_YAML_NAME;
use a8e_core::config::extensions::get_enabled_extensions_with_config;
use a8e_core::config::paths::Paths;
use a8e_core::config::permission::PermissionManager;
use a8e_core::config::Config;
use a8e_core::conversation::message::{ActionRequiredData, Message, MessageContent};
use a8e_core::conversation::Conversation;
use a8e_core::mcp_utils::ToolResult;
use a8e_core::permission::permission_confirmation::PrincipalType;
use a8e_core::permission::{Permission, PermissionConfirmation};
use a8e_core::providers::base::Provider;
use a8e_core::providers::provider_registry::ProviderConstructor;
use a8e_core::session::session_manager::SessionType;
use a8e_core::session::{Session, SessionManager};
use anyhow::Result;
use fs_err as fs;
use rmcp::model::{CallToolResult, RawContent, ResourceContents, Role};
use sacp::schema::{
AgentCapabilities, AuthMethod, AuthenticateRequest, AuthenticateResponse, BlobResourceContents,
CancelNotification, Content, ContentBlock, ContentChunk, EmbeddedResource,
EmbeddedResourceResource, ImageContent, InitializeRequest, InitializeResponse,
LoadSessionRequest, LoadSessionResponse, McpCapabilities, McpServer, ModelId, ModelInfo,
NewSessionRequest, NewSessionResponse, PermissionOption, PermissionOptionKind,
PromptCapabilities, PromptRequest, PromptResponse, RequestPermissionOutcome,
RequestPermissionRequest, ResourceLink, SessionId, SessionModelState, SessionNotification,
SessionUpdate, SetSessionModelRequest, SetSessionModelResponse, StopReason, TextContent,
TextResourceContents, ToolCall, ToolCallContent, ToolCallId, ToolCallLocation, ToolCallStatus,
ToolCallUpdate, ToolCallUpdateFields, ToolKind,
};
use sacp::{AgentToClient, ByteStreams, Handled, JrConnectionCx, JrMessageHandler, MessageCx};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _};
use tokio_util::sync::CancellationToken;
use tracing::{debug, error, info, warn};
use url::Url;
struct A8eAcpSession {
agent: Arc<Agent>,
messages: Conversation,
tool_requests: HashMap<String, a8e_core::conversation::message::ToolRequest>,
cancel_token: Option<CancellationToken>,
}
pub struct A8eAcpAgent {
sessions: Arc<Mutex<HashMap<String, A8eAcpSession>>>,
provider_factory: ProviderConstructor,
config_dir: std::path::PathBuf,
session_manager: Arc<SessionManager>,
permission_manager: Arc<PermissionManager>,
a8e_mode: a8e_core::config::A8eMode,
disable_session_naming: bool,
builtins: Vec<String>,
}
fn mcp_server_to_extension_config(mcp_server: McpServer) -> Result<ExtensionConfig, String> {
match mcp_server {
McpServer::Stdio(stdio) => Ok(ExtensionConfig::Stdio {
name: stdio.name,
description: String::new(),
cmd: stdio.command.to_string_lossy().to_string(),
args: stdio.args,
envs: Envs::new(stdio.env.into_iter().map(|e| (e.name, e.value)).collect()),
env_keys: vec![],
timeout: None,
bundled: Some(false),
available_tools: vec![],
}),
McpServer::Http(http) => Ok(ExtensionConfig::StreamableHttp {
name: http.name,
description: String::new(),
uri: http.url,
envs: Envs::default(),
env_keys: vec![],
headers: http
.headers
.into_iter()
.map(|h| (h.name, h.value))
.collect(),
timeout: None,
bundled: Some(false),
available_tools: vec![],
}),
McpServer::Sse(_) => Err("SSE is unsupported, migrate to streamable_http".to_string()),
_ => Err("Unknown MCP server type".to_string()),
}
}
fn create_tool_location(path: &str, line: Option<u32>) -> ToolCallLocation {
let mut loc = ToolCallLocation::new(path);
if let Some(l) = line {
loc = loc.line(l);
}
loc
}
fn extract_tool_locations(
tool_request: &a8e_core::conversation::message::ToolRequest,
tool_response: &a8e_core::conversation::message::ToolResponse,
) -> Vec<ToolCallLocation> {
let mut locations = Vec::new();
if let Ok(tool_call) = &tool_request.tool_call {
if tool_call.name != "developer__text_editor" {
return locations;
}
let path_str = tool_call
.arguments
.as_ref()
.and_then(|args| args.get("path"))
.and_then(|p| p.as_str());
if let Some(path_str) = path_str {
let command = tool_call
.arguments
.as_ref()
.and_then(|args| args.get("command"))
.and_then(|c| c.as_str());
if let Ok(result) = &tool_response.tool_result {
for content in &result.content {
if let RawContent::Text(text_content) = &content.raw {
let text = &text_content.text;
match command {
Some("view") => {
let line = extract_view_line_range(text)
.map(|range| range.0 as u32)
.or(Some(1));
locations.push(create_tool_location(path_str, line));
}
Some("str_replace") | Some("insert") => {
let line = extract_first_line_number(text)
.map(|l| l as u32)
.or(Some(1));
locations.push(create_tool_location(path_str, line));
}
Some("write") => {
locations.push(create_tool_location(path_str, Some(1)));
}
_ => {
locations.push(create_tool_location(path_str, Some(1)));
}
}
break;
}
}
}
if locations.is_empty() {
locations.push(create_tool_location(path_str, Some(1)));
}
}
}
locations
}
fn extract_view_line_range(text: &str) -> Option<(usize, usize)> {
let re = regex::Regex::new(r"\(lines (\d+)-(\d+|end)\)").ok()?;
if let Some(caps) = re.captures(text) {
let start = caps.get(1)?.as_str().parse::<usize>().ok()?;
let end = if caps.get(2)?.as_str() == "end" {
start
} else {
caps.get(2)?.as_str().parse::<usize>().ok()?
};
return Some((start, end));
}
None
}
fn extract_first_line_number(text: &str) -> Option<usize> {
let re = regex::Regex::new(r"```[^\n]*\n(\d+):").ok()?;
if let Some(caps) = re.captures(text) {
return caps.get(1)?.as_str().parse::<usize>().ok();
}
None
}
fn read_resource_link(link: ResourceLink) -> Option<String> {
let url = Url::parse(&link.uri).ok()?;
if url.scheme() == "file" {
let path = url.to_file_path().ok()?;
let contents = fs::read_to_string(&path).ok()?;
Some(format!(
"\n\n# {}\n```\n{}\n```",
path.to_string_lossy(),
contents
))
} else {
None
}
}
fn format_tool_name(tool_name: &str) -> String {
let capitalize = |s: &str| {
s.split_whitespace()
.map(|word| {
let mut chars = word.chars();
match chars.next() {
None => String::new(),
Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
}
})
.collect::<Vec<_>>()
.join(" ")
};
if let Some((extension, tool)) = tool_name.split_once("__") {
let formatted_extension = extension.replace('_', " ");
let formatted_tool = tool.replace('_', " ");
format!(
"{}: {}",
capitalize(&formatted_extension),
capitalize(&formatted_tool)
)
} else {
let formatted = tool_name.replace('_', " ");
capitalize(&formatted)
}
}
async fn add_builtins(agent: &Agent, builtins: Vec<String>) {
for builtin in builtins {
let config = if PLATFORM_EXTENSIONS.contains_key(builtin.as_str()) {
ExtensionConfig::Platform {
name: builtin.clone(),
description: builtin.clone(),
display_name: None,
bundled: None,
available_tools: Vec::new(),
}
} else {
ExtensionConfig::Builtin {
name: builtin.clone(),
display_name: None,
timeout: None,
bundled: None,
description: builtin.clone(),
available_tools: Vec::new(),
}
};
match agent
.extension_manager
.add_extension(config, None, None, None)
.await
{
Ok(_) => info!(extension = %builtin, "extension loaded"),
Err(e) => warn!(extension = %builtin, error = %e, "extension load failed"),
}
}
}
async fn add_extensions(agent: &Agent, extensions: Vec<ExtensionConfig>) {
for extension in extensions {
let name = extension.name().to_string();
match agent
.extension_manager
.add_extension(extension, None, None, None)
.await
{
Ok(_) => info!(extension = %name, "extension loaded"),
Err(e) => warn!(extension = %name, error = %e, "extension load failed"),
}
}
}
async fn build_model_state(
provider: &dyn Provider,
current_model: &str,
) -> Result<SessionModelState, sacp::Error> {
let models = provider.fetch_recommended_models().await.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to fetch models: {}", e))
})?;
Ok(SessionModelState::new(
ModelId::new(current_model),
models
.iter()
.map(|name| ModelInfo::new(ModelId::new(&**name), &**name))
.collect(),
))
}
impl A8eAcpAgent {
pub fn permission_manager(&self) -> Arc<PermissionManager> {
Arc::clone(&self.permission_manager)
}
pub async fn new(
provider_factory: ProviderConstructor,
builtins: Vec<String>,
data_dir: std::path::PathBuf,
config_dir: std::path::PathBuf,
a8e_mode: a8e_core::config::A8eMode,
disable_session_naming: bool,
) -> Result<Self> {
let session_manager = Arc::new(SessionManager::new(data_dir));
let permission_manager = Arc::new(PermissionManager::new(config_dir.clone()));
Ok(Self {
sessions: Arc::new(Mutex::new(HashMap::new())),
provider_factory,
config_dir,
session_manager,
permission_manager,
a8e_mode,
disable_session_naming,
builtins,
})
}
async fn create_agent_for_session(&self) -> Arc<Agent> {
let agent = Agent::with_config(AgentConfig::new(
Arc::clone(&self.session_manager),
Arc::clone(&self.permission_manager),
None,
self.a8e_mode,
self.disable_session_naming,
A8ePlatform::Cli,
));
let agent = Arc::new(agent);
let config_path = self.config_dir.join(CONFIG_YAML_NAME);
if let Ok(config_file) = Config::new(&config_path, "a8e") {
let extensions = get_enabled_extensions_with_config(&config_file);
add_extensions(&agent, extensions).await;
}
add_builtins(&agent, self.builtins.clone()).await;
agent
}
pub async fn has_session(&self, session_id: &str) -> bool {
self.sessions.lock().await.contains_key(session_id)
}
fn convert_acp_prompt_to_message(&self, prompt: Vec<ContentBlock>) -> Message {
let mut user_message = Message::user();
for block in prompt {
match block {
ContentBlock::Text(text) => {
user_message = user_message.with_text(&text.text);
}
ContentBlock::Image(image) => {
user_message = user_message.with_image(&image.data, &image.mime_type);
}
ContentBlock::Resource(resource) => {
if let EmbeddedResourceResource::TextResourceContents(text_resource) =
&resource.resource
{
let header = format!("--- Resource: {} ---\n", text_resource.uri);
let content = format!("{}{}\n---\n", header, text_resource.text);
user_message = user_message.with_text(&content);
}
}
ContentBlock::ResourceLink(link) => {
if let Some(text) = read_resource_link(link) {
user_message = user_message.with_text(text)
}
}
ContentBlock::Audio(..) | _ => (),
}
}
user_message
}
async fn handle_message_content(
&self,
content_item: &MessageContent,
session_id: &SessionId,
session: &mut A8eAcpSession,
cx: &JrConnectionCx<AgentToClient>,
) -> Result<(), sacp::Error> {
match content_item {
MessageContent::Text(text) => {
cx.send_notification(SessionNotification::new(
session_id.clone(),
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new(text.text.clone()),
))),
))?;
}
MessageContent::ToolRequest(tool_request) => {
self.handle_tool_request(tool_request, session_id, session, cx)
.await?;
}
MessageContent::ToolResponse(tool_response) => {
self.handle_tool_response(tool_response, session_id, session, cx)
.await?;
}
MessageContent::Thinking(thinking) => {
cx.send_notification(SessionNotification::new(
session_id.clone(),
SessionUpdate::AgentThoughtChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new(thinking.thinking.clone()),
))),
))?;
}
MessageContent::ActionRequired(action_required) => {
if let ActionRequiredData::ToolConfirmation {
id,
tool_name,
arguments,
prompt,
} = &action_required.data
{
self.handle_tool_permission_request(
cx,
&session.agent,
session_id,
id.clone(),
tool_name.clone(),
arguments.clone(),
prompt.clone(),
)?;
}
}
_ => {}
}
Ok(())
}
async fn handle_tool_request(
&self,
tool_request: &a8e_core::conversation::message::ToolRequest,
session_id: &SessionId,
session: &mut A8eAcpSession,
cx: &JrConnectionCx<AgentToClient>,
) -> Result<(), sacp::Error> {
session
.tool_requests
.insert(tool_request.id.clone(), tool_request.clone());
let tool_name = match &tool_request.tool_call {
Ok(tool_call) => tool_call.name.to_string(),
Err(_) => "error".to_string(),
};
cx.send_notification(SessionNotification::new(
session_id.clone(),
SessionUpdate::ToolCall(
ToolCall::new(
ToolCallId::new(tool_request.id.clone()),
format_tool_name(&tool_name),
)
.status(ToolCallStatus::Pending),
),
))?;
Ok(())
}
async fn handle_tool_response(
&self,
tool_response: &a8e_core::conversation::message::ToolResponse,
session_id: &SessionId,
session: &mut A8eAcpSession,
cx: &JrConnectionCx<AgentToClient>,
) -> Result<(), sacp::Error> {
let status = match &tool_response.tool_result {
Ok(result) if result.is_error == Some(true) => ToolCallStatus::Failed,
Ok(_) => ToolCallStatus::Completed,
Err(_) => ToolCallStatus::Failed,
};
let content = build_tool_call_content(&tool_response.tool_result);
let locations = if let Some(tool_request) = session.tool_requests.get(&tool_response.id) {
extract_tool_locations(tool_request, tool_response)
} else {
Vec::new()
};
let mut fields = ToolCallUpdateFields::new().status(status).content(content);
if !locations.is_empty() {
fields = fields.locations(locations);
}
cx.send_notification(SessionNotification::new(
session_id.clone(),
SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(
ToolCallId::new(tool_response.id.clone()),
fields,
)),
))?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn handle_tool_permission_request(
&self,
cx: &JrConnectionCx<AgentToClient>,
agent: &Arc<Agent>,
session_id: &SessionId,
request_id: String,
tool_name: String,
arguments: serde_json::Map<String, serde_json::Value>,
prompt: Option<String>,
) -> Result<(), sacp::Error> {
let cx = cx.clone();
let agent = agent.clone();
let session_id = session_id.clone();
let formatted_name = format_tool_name(&tool_name);
let mut fields = ToolCallUpdateFields::new()
.title(formatted_name)
.kind(ToolKind::default())
.status(ToolCallStatus::Pending)
.raw_input(serde_json::Value::Object(arguments));
if let Some(p) = prompt {
fields = fields.content(vec![ToolCallContent::Content(Content::new(
ContentBlock::Text(TextContent::new(p)),
))]);
}
let tool_call_update = ToolCallUpdate::new(ToolCallId::new(request_id.clone()), fields);
fn option(kind: PermissionOptionKind) -> PermissionOption {
let id = serde_json::to_value(kind)
.unwrap()
.as_str()
.unwrap()
.to_string();
PermissionOption::new(id.clone(), id, kind)
}
let options = vec![
option(PermissionOptionKind::AllowAlways),
option(PermissionOptionKind::AllowOnce),
option(PermissionOptionKind::RejectOnce),
option(PermissionOptionKind::RejectAlways),
];
let permission_request =
RequestPermissionRequest::new(session_id, tool_call_update, options);
cx.send_request(permission_request)
.on_receiving_result(move |result| async move {
match result {
Ok(response) => {
agent
.handle_confirmation(
request_id,
outcome_to_confirmation(&response.outcome),
)
.await;
Ok(())
}
Err(e) => {
error!(error = ?e, "permission request failed");
agent
.handle_confirmation(
request_id,
PermissionConfirmation {
principal_type: PrincipalType::Tool,
permission: Permission::Cancel,
},
)
.await;
Ok(())
}
}
})?;
Ok(())
}
}
fn outcome_to_confirmation(outcome: &RequestPermissionOutcome) -> PermissionConfirmation {
let permission = match outcome {
RequestPermissionOutcome::Cancelled => Permission::Cancel,
RequestPermissionOutcome::Selected(selected) => {
match serde_json::from_value::<PermissionOptionKind>(serde_json::Value::String(
selected.option_id.0.to_string(),
)) {
Ok(PermissionOptionKind::AllowAlways) => Permission::AlwaysAllow,
Ok(PermissionOptionKind::AllowOnce) => Permission::AllowOnce,
Ok(PermissionOptionKind::RejectOnce) => Permission::DenyOnce,
Ok(PermissionOptionKind::RejectAlways) => Permission::AlwaysDeny,
_ => Permission::Cancel,
}
}
_ => Permission::Cancel,
};
PermissionConfirmation {
principal_type: PrincipalType::Tool,
permission,
}
}
fn build_tool_call_content(tool_result: &ToolResult<CallToolResult>) -> Vec<ToolCallContent> {
match tool_result {
Ok(result) => result
.content
.iter()
.filter_map(|content| match &content.raw {
RawContent::Text(val) => Some(ToolCallContent::Content(Content::new(
ContentBlock::Text(TextContent::new(val.text.clone())),
))),
RawContent::Image(val) => Some(ToolCallContent::Content(Content::new(
ContentBlock::Image(ImageContent::new(val.data.clone(), val.mime_type.clone())),
))),
RawContent::Resource(val) => {
let resource = match &val.resource {
ResourceContents::TextResourceContents {
mime_type,
text,
uri,
..
} => EmbeddedResourceResource::TextResourceContents(
TextResourceContents::new(text.clone(), uri.clone())
.mime_type(mime_type.clone()),
),
ResourceContents::BlobResourceContents {
mime_type,
blob,
uri,
..
} => EmbeddedResourceResource::BlobResourceContents(
BlobResourceContents::new(blob.clone(), uri.clone())
.mime_type(mime_type.clone()),
),
};
Some(ToolCallContent::Content(Content::new(
ContentBlock::Resource(EmbeddedResource::new(resource)),
)))
}
RawContent::Audio(_) | RawContent::ResourceLink(_) => None,
})
.collect(),
Err(_) => Vec::new(),
}
}
impl A8eAcpAgent {
async fn on_initialize(
&self,
args: InitializeRequest,
) -> Result<InitializeResponse, sacp::Error> {
debug!(?args, "initialize request");
let capabilities = AgentCapabilities::new()
.load_session(true)
.prompt_capabilities(
PromptCapabilities::new()
.image(true)
.audio(false)
.embedded_context(true),
)
.mcp_capabilities(McpCapabilities::new().http(true));
Ok(InitializeResponse::new(args.protocol_version)
.agent_capabilities(capabilities)
.auth_methods(vec![AuthMethod::new("a8e-provider", "Configure Provider")
.description(
"Run `a8e configure` to set up your AI provider and API key",
)]))
}
async fn on_new_session(
&self,
args: NewSessionRequest,
) -> Result<NewSessionResponse, sacp::Error> {
debug!(?args, "new session request");
let a8e_session = self
.session_manager
.create_session(
args.cwd.clone(),
"ACP Session".to_string(),
SessionType::User,
)
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to create session: {}", e))
})?;
let agent = self.create_agent_for_session().await;
let provider = self
.init_provider(&agent, &a8e_session)
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to set provider: {}", e))
})?;
for mcp_server in args.mcp_servers {
let config = match mcp_server_to_extension_config(mcp_server) {
Ok(c) => c,
Err(msg) => {
return Err(sacp::Error::invalid_params().data(msg));
}
};
let name = config.name().to_string();
if let Err(e) = agent.add_extension(config, &a8e_session.id).await {
return Err(sacp::Error::internal_error()
.data(format!("Failed to add MCP server '{}': {}", name, e)));
}
}
let session = A8eAcpSession {
agent,
messages: Conversation::new_unvalidated(Vec::new()),
tool_requests: HashMap::new(),
cancel_token: None,
};
let mut sessions = self.sessions.lock().await;
sessions.insert(a8e_session.id.clone(), session);
info!(
session_id = %a8e_session.id,
session_type = "acp",
"Session started"
);
let model_state =
build_model_state(&*provider, &provider.get_model_config().model_name).await?;
Ok(NewSessionResponse::new(SessionId::new(a8e_session.id)).models(model_state))
}
async fn init_provider(&self, agent: &Agent, session: &Session) -> Result<Arc<dyn Provider>> {
let model_config = match &session.model_config {
Some(config) => config.clone(),
None => {
let config_path = self.config_dir.join(CONFIG_YAML_NAME);
let config = Config::new(&config_path, "a8e")?;
let model_id = config.get_a8e_model()?;
let provider_name = config.get_a8e_provider()?;
a8e_core::model::ModelConfig::new(&model_id)?.with_canonical_limits(&provider_name)
}
};
let provider = (self.provider_factory)(model_config, Vec::new()).await?;
agent.update_provider(provider.clone(), &session.id).await?;
Ok(provider)
}
async fn on_load_session(
&self,
args: LoadSessionRequest,
cx: &JrConnectionCx<AgentToClient>,
) -> Result<LoadSessionResponse, sacp::Error> {
debug!(?args, "load session request");
let session_id = args.session_id.0.to_string();
let a8e_session = self
.session_manager
.get_session(&session_id, true)
.await
.map_err(|e| {
sacp::Error::invalid_params()
.data(format!("Failed to load session {}: {}", session_id, e))
})?;
let agent = self.create_agent_for_session().await;
let provider = self
.init_provider(&agent, &a8e_session)
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to set provider: {}", e))
})?;
let conversation = a8e_session.conversation.ok_or_else(|| {
sacp::Error::internal_error()
.data(format!("Session {} has no conversation data", session_id))
})?;
self.session_manager
.update(&session_id)
.working_dir(args.cwd.clone())
.apply()
.await
.map_err(|e| {
sacp::Error::internal_error()
.data(format!("Failed to update session working directory: {}", e))
})?;
let mut session = A8eAcpSession {
agent,
messages: conversation.clone(),
tool_requests: HashMap::new(),
cancel_token: None,
};
for message in conversation.messages() {
if !message.metadata.user_visible {
continue;
}
for content_item in &message.content {
match content_item {
MessageContent::Text(text) => {
let chunk = ContentChunk::new(ContentBlock::Text(TextContent::new(
text.text.clone(),
)));
let update = match message.role {
Role::User => SessionUpdate::UserMessageChunk(chunk),
Role::Assistant => SessionUpdate::AgentMessageChunk(chunk),
};
cx.send_notification(SessionNotification::new(
args.session_id.clone(),
update,
))?;
}
MessageContent::ToolRequest(tool_request) => {
self.handle_tool_request(tool_request, &args.session_id, &mut session, cx)
.await?;
}
MessageContent::ToolResponse(tool_response) => {
self.handle_tool_response(
tool_response,
&args.session_id,
&mut session,
cx,
)
.await?;
}
MessageContent::Thinking(thinking) => {
cx.send_notification(SessionNotification::new(
args.session_id.clone(),
SessionUpdate::AgentThoughtChunk(ContentChunk::new(
ContentBlock::Text(TextContent::new(thinking.thinking.clone())),
)),
))?;
}
_ => {}
}
}
}
let mut sessions = self.sessions.lock().await;
sessions.insert(session_id.clone(), session);
info!(
session_id = %session_id,
session_type = "acp",
"Session loaded"
);
let model_state =
build_model_state(&*provider, &provider.get_model_config().model_name).await?;
Ok(LoadSessionResponse::new().models(model_state))
}
async fn on_prompt(
&self,
args: PromptRequest,
cx: &JrConnectionCx<AgentToClient>,
) -> Result<PromptResponse, sacp::Error> {
let session_id = args.session_id.0.to_string();
let cancel_token = CancellationToken::new();
let agent = {
let mut sessions = self.sessions.lock().await;
let session = sessions.get_mut(&session_id).ok_or_else(|| {
sacp::Error::invalid_params().data(format!("Session not found: {}", session_id))
})?;
session.cancel_token = Some(cancel_token.clone());
session.agent.clone()
};
let user_message = self.convert_acp_prompt_to_message(args.prompt);
let session_config = SessionConfig {
id: session_id.clone(),
schedule_id: None,
max_turns: None,
retry_config: None,
};
let mut stream = agent
.reply(user_message, session_config, Some(cancel_token.clone()))
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Error getting agent reply: {}", e))
})?;
use futures::StreamExt;
let mut was_cancelled = false;
while let Some(event) = stream.next().await {
if cancel_token.is_cancelled() {
was_cancelled = true;
break;
}
match event {
Ok(a8e_core::agents::AgentEvent::Message(message)) => {
let mut sessions = self.sessions.lock().await;
let session = sessions.get_mut(&session_id).ok_or_else(|| {
sacp::Error::invalid_params()
.data(format!("Session not found: {}", session_id))
})?;
session.messages.push(message.clone());
for content_item in &message.content {
self.handle_message_content(content_item, &args.session_id, session, cx)
.await?;
}
}
Ok(_) => {}
Err(e) => {
return Err(sacp::Error::internal_error()
.data(format!("Error in agent response stream: {}", e)));
}
}
}
let mut sessions = self.sessions.lock().await;
if let Some(session) = sessions.get_mut(&session_id) {
session.cancel_token = None;
}
Ok(PromptResponse::new(if was_cancelled {
StopReason::Cancelled
} else {
StopReason::EndTurn
}))
}
async fn on_cancel(&self, args: CancelNotification) -> Result<(), sacp::Error> {
debug!(?args, "cancel request");
let session_id = args.session_id.0.to_string();
let mut sessions = self.sessions.lock().await;
if let Some(session) = sessions.get_mut(&session_id) {
if let Some(ref token) = session.cancel_token {
info!(session_id = %session_id, "prompt cancelled");
token.cancel();
}
} else {
warn!(session_id = %session_id, "cancel request for unknown session");
}
Ok(())
}
async fn on_set_model(
&self,
session_id: &str,
model_id: &str,
) -> Result<SetSessionModelResponse, sacp::Error> {
let config_path = self.config_dir.join(CONFIG_YAML_NAME);
let config = Config::new(&config_path, "a8e").map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to read config: {}", e))
})?;
let provider_name = config.get_a8e_provider().map_err(|_| {
sacp::Error::internal_error().data("No provider configured".to_string())
})?;
let model_config = a8e_core::model::ModelConfig::new(model_id)
.map_err(|e| {
sacp::Error::invalid_params().data(format!("Invalid model config: {}", e))
})?
.with_canonical_limits(&provider_name);
let provider = (self.provider_factory)(model_config, Vec::new())
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to create provider: {}", e))
})?;
let agent = {
let sessions = self.sessions.lock().await;
let session = sessions.get(session_id).ok_or_else(|| {
sacp::Error::invalid_params().data(format!("Session not found: {}", session_id))
})?;
session.agent.clone()
};
agent
.update_provider(provider, session_id)
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to update provider: {}", e))
})?;
info!(session_id = %session_id, model_id = %model_id, "Model switched");
Ok(SetSessionModelResponse::new())
}
}
pub struct A8eAcpHandler {
pub agent: Arc<A8eAcpAgent>,
}
impl JrMessageHandler for A8eAcpHandler {
type Link = AgentToClient;
fn describe_chain(&self) -> impl std::fmt::Debug {
"a8e-acp"
}
async fn handle_message(
&mut self,
message: MessageCx,
cx: JrConnectionCx<AgentToClient>,
) -> Result<Handled<MessageCx>, sacp::Error> {
use sacp::util::MatchMessageFrom;
use sacp::JrRequestCx;
MatchMessageFrom::new(message, &cx)
.if_request(
|req: InitializeRequest, req_cx: JrRequestCx<InitializeResponse>| async {
req_cx.respond(self.agent.on_initialize(req).await?)
},
)
.await
.if_request(
|_req: AuthenticateRequest, req_cx: JrRequestCx<AuthenticateResponse>| async {
req_cx.respond(AuthenticateResponse::new())
},
)
.await
.if_request(
|req: NewSessionRequest, req_cx: JrRequestCx<NewSessionResponse>| async {
req_cx.respond(self.agent.on_new_session(req).await?)
},
)
.await
.if_request(
|req: LoadSessionRequest, req_cx: JrRequestCx<LoadSessionResponse>| async {
req_cx.respond(self.agent.on_load_session(req, &cx).await?)
},
)
.await
.if_request(
|req: PromptRequest, req_cx: JrRequestCx<PromptResponse>| async {
let agent = self.agent.clone();
let cx_clone = cx.clone();
cx.spawn(async move {
match agent.on_prompt(req, &cx_clone).await {
Ok(response) => {
req_cx.respond(response)?;
}
Err(e) => {
req_cx.respond_with_error(e)?;
}
}
Ok(())
})?;
Ok(())
},
)
.await
.if_notification(|notif: CancelNotification| async {
self.agent.on_cancel(notif).await
})
.await
.otherwise({
let agent = self.agent.clone();
|message: MessageCx| async move {
match message {
MessageCx::Request(req, request_cx)
if req.method == "session/set_model" =>
{
let params: SetSessionModelRequest = serde_json::from_value(req.params)
.map_err(|e| sacp::Error::invalid_params().data(e.to_string()))?;
let resp = agent
.on_set_model(¶ms.session_id.0, ¶ms.model_id.0)
.await?;
let json = serde_json::to_value(resp)
.map_err(|e| sacp::Error::internal_error().data(e.to_string()))?;
request_cx.respond(json)?;
Ok(())
}
_ => Err(sacp::Error::method_not_found()),
}
}
})
.await
.map(|()| Handled::Yes)
}
}
pub async fn serve<R, W>(agent: Arc<A8eAcpAgent>, read: R, write: W) -> Result<()>
where
R: futures::AsyncRead + Unpin + Send + 'static,
W: futures::AsyncWrite + Unpin + Send + 'static,
{
let handler = A8eAcpHandler { agent };
AgentToClient::builder()
.name("a8e-acp")
.with_handler(handler)
.serve(ByteStreams::new(write, read))
.await?;
Ok(())
}
pub async fn run(builtins: Vec<String>) -> Result<()> {
register_builtin_extensions(a8e_mcp::BUILTIN_EXTENSIONS.clone());
info!("listening on stdio");
let outgoing = tokio::io::stdout().compat_write();
let incoming = tokio::io::stdin().compat();
let server =
crate::server_factory::AcpServer::new(crate::server_factory::AcpServerFactoryConfig {
builtins,
data_dir: Paths::data_dir(),
config_dir: Paths::config_dir(),
});
let agent = server.create_agent().await?;
serve(agent, incoming, outgoing).await
}
#[cfg(test)]
mod tests {
use super::*;
use sacp::schema::{
EnvVariable, HttpHeader, McpServer, McpServerHttp, McpServerSse, McpServerStdio,
PermissionOptionId, ResourceLink, SelectedPermissionOutcome,
};
use std::io::Write;
use tempfile::NamedTempFile;
use test_case::test_case;
#[test_case(
McpServer::Stdio(
McpServerStdio::new("github", "/path/to/github-mcp-server")
.args(vec!["stdio".into()])
.env(vec![EnvVariable::new("GITHUB_PERSONAL_ACCESS_TOKEN", "ghp_xxxxxxxxxxxx")])
),
Ok(ExtensionConfig::Stdio {
name: "github".into(),
description: String::new(),
cmd: "/path/to/github-mcp-server".into(),
args: vec!["stdio".into()],
envs: Envs::new(
[(
"GITHUB_PERSONAL_ACCESS_TOKEN".into(),
"ghp_xxxxxxxxxxxx".into()
)]
.into()
),
env_keys: vec![],
timeout: None,
bundled: Some(false),
available_tools: vec![],
})
)]
#[test_case(
McpServer::Http(
McpServerHttp::new("github", "https://api.githubcopilot.com/mcp/")
.headers(vec![HttpHeader::new("Authorization", "Bearer ghp_xxxxxxxxxxxx")])
),
Ok(ExtensionConfig::StreamableHttp {
name: "github".into(),
description: String::new(),
uri: "https://api.githubcopilot.com/mcp/".into(),
envs: Envs::default(),
env_keys: vec![],
headers: HashMap::from([(
"Authorization".into(),
"Bearer ghp_xxxxxxxxxxxx".into()
)]),
timeout: None,
bundled: Some(false),
available_tools: vec![],
})
)]
#[test_case(
McpServer::Sse(McpServerSse::new("test-sse", "https://agent-fin.biodnd.com/sse")),
Err("SSE is unsupported, migrate to streamable_http".to_string())
)]
fn test_mcp_server_to_extension_config(
input: McpServer,
expected: Result<ExtensionConfig, String>,
) {
assert_eq!(mcp_server_to_extension_config(input), expected);
}
fn new_resource_link(content: &str) -> anyhow::Result<(ResourceLink, NamedTempFile)> {
let mut file = NamedTempFile::new()?;
file.write_all(content.as_bytes())?;
let name = file
.path()
.file_name()
.unwrap()
.to_string_lossy()
.to_string();
let uri = format!("file://{}", file.path().to_str().unwrap());
let link = ResourceLink::new(name, uri);
Ok((link, file))
}
#[test]
fn test_read_resource_link_non_file_scheme() {
let (link, file) = new_resource_link("print(\"hello, world\")").unwrap();
let result = read_resource_link(link).unwrap();
let expected = format!(
"
# {}
```
print(\"hello, world\")
```",
file.path().to_str().unwrap(),
);
assert_eq!(result, expected,)
}
#[test]
fn test_format_tool_name_with_extension() {
assert_eq!(
format_tool_name("developer__text_editor"),
"Developer: Text Editor"
);
assert_eq!(
format_tool_name("platform__manage_extensions"),
"Platform: Manage Extensions"
);
assert_eq!(format_tool_name("todo__write"), "Todo: Write");
}
#[test]
fn test_format_tool_name_without_extension() {
assert_eq!(format_tool_name("simple_tool"), "Simple Tool");
assert_eq!(format_tool_name("another_name"), "Another Name");
assert_eq!(format_tool_name("single"), "Single");
}
#[test_case(
RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(PermissionOptionId::from("allow_once".to_string()))),
PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::AllowOnce };
"allow_once_maps_to_allow_once"
)]
#[test_case(
RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(PermissionOptionId::from("allow_always".to_string()))),
PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::AlwaysAllow };
"allow_always_maps_to_always_allow"
)]
#[test_case(
RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(PermissionOptionId::from("reject_once".to_string()))),
PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::DenyOnce };
"reject_once_maps_to_deny_once"
)]
#[test_case(
RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(PermissionOptionId::from("reject_always".to_string()))),
PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::AlwaysDeny };
"reject_always_maps_to_always_deny"
)]
#[test_case(
RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(PermissionOptionId::from("unknown".to_string()))),
PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::Cancel };
"unknown_option_maps_to_cancel"
)]
#[test_case(
RequestPermissionOutcome::Cancelled,
PermissionConfirmation { principal_type: PrincipalType::Tool, permission: Permission::Cancel };
"cancelled_maps_to_cancel"
)]
fn test_outcome_to_confirmation(
input: RequestPermissionOutcome,
expected: PermissionConfirmation,
) {
assert_eq!(outcome_to_confirmation(&input), expected);
}
use a8e_core::providers::errors::ProviderError;
struct MockModelProvider {
models: Result<Vec<String>, ProviderError>,
}
#[async_trait::async_trait]
impl a8e_core::providers::base::Provider for MockModelProvider {
fn get_name(&self) -> &str {
"mock"
}
async fn stream(
&self,
_model_config: &a8e_core::model::ModelConfig,
_session_id: &str,
_system: &str,
_messages: &[a8e_core::conversation::message::Message],
_tools: &[rmcp::model::Tool],
) -> Result<a8e_core::providers::base::MessageStream, ProviderError> {
unimplemented!()
}
fn get_model_config(&self) -> a8e_core::model::ModelConfig {
a8e_core::model::ModelConfig::new_or_fail("unused")
}
async fn fetch_recommended_models(&self) -> Result<Vec<String>, ProviderError> {
self.models.clone()
}
}
#[test_case(
"model-a", Ok(vec!["model-a".into(), "model-b".into()])
=> Ok(SessionModelState::new(
ModelId::new("model-a"),
vec![ModelInfo::new(ModelId::new("model-a"), "model-a"),
ModelInfo::new(ModelId::new("model-b"), "model-b")],
))
; "returns current and available models"
)]
#[test_case(
"model-a", Ok(vec![])
=> Ok(SessionModelState::new(ModelId::new("model-a"), vec![]))
; "empty model list"
)]
#[test_case(
"model-a", Err(ProviderError::ExecutionError("fail".into()))
=> matches Err(_)
; "fetch error propagates"
)]
#[test_case(
"switched-model", Ok(vec!["model-a".into(), "switched-model".into()])
=> Ok(SessionModelState::new(
ModelId::new("switched-model"),
vec![ModelInfo::new(ModelId::new("model-a"), "model-a"),
ModelInfo::new(ModelId::new("switched-model"), "switched-model")],
))
; "current model reflects switched model"
)]
#[tokio::test]
async fn test_build_model_state(
current_model: &str,
models: Result<Vec<String>, ProviderError>,
) -> Result<SessionModelState, sacp::Error> {
let provider = MockModelProvider { models };
build_model_state(&provider, current_model).await
}
}