refactor: 迁移 Rust 核心代码至 rust/ 目录

- 将 Rust 核心从 astrbot/rust/ 迁移至 rust/
- 新增 a2a.rs: Agent-to-Agent 通信协议
- 新增 abp.rs: ABP 插件协议客户端
- 新增 server.rs: WebSocket/HTTP 服务器
- 更新 Cargo.toml 依赖 (futures-util)
- 重构 config.rs, orchestrator.rs, protocol.rs 等核心模块
This commit is contained in:
LIghtJUNction
2026-03-27 00:15:44 +08:00
parent e1ac6733e8
commit eae35ffd64
13 changed files with 5781 additions and 382 deletions

16
astrbot/rust/_core.pyi Normal file
View File

@@ -0,0 +1,16 @@
from typing import Any
class AstrbotOrchestrator:
def start(self) -> None: ...
def stop(self) -> None: ...
def is_running(self) -> bool: ...
def register_star(self, name: str, handler: str) -> None: ...
def unregister_star(self, name: str) -> None: ...
def list_stars(self) -> list[str]: ...
def record_activity(self) -> None: ...
def get_stats(self) -> dict[str, Any]: ...
def set_protocol_connected(self, protocol: str, connected: bool) -> None: ...
def get_protocol_status(self, protocol: str) -> dict[str, Any] | None: ...
def get_orchestrator() -> AstrbotOrchestrator: ...
def cli(args: list[str]) -> None: ...

1850
rust/Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@@ -4,12 +4,14 @@ version = "4.25.0"
edition = "2024"
[lib]
name = "astrbot__core"
name = "astrbot_core"
crate-type = ["cdylib", "rlib"]
[features]
default = ["python"]
python = ["pyo3"]
python = ["pyo3", "pyo3-asyncio"]
pyo3-asyncio = []
pyo3-async-runtimes = ["dep:pyo3-async-runtimes"]
[dependencies]
serde = { version = "1", features = ["derive"] }
@@ -27,6 +29,13 @@ pyo3 = { version = "0.28", features = [
"full",
"extension-module",
], optional = true }
pyo3-async-runtimes = { version = "0.28", optional = true }
uuid = { version = "1.22.0", features = ["v4"] }
reqwest = { version = "0.13.2", features = ["json", "default"] }
dirs = "6.0.0"
tokio-tungstenite = "0.29.0"
base64 = "0.22.1"
futures-util = "0.3"
[profile.release]
opt-level = 3

464
rust/src/a2a.rs Normal file
View File

@@ -0,0 +1,464 @@
//! A2A (Agent-to-Agent) Protocol Implementation
//!
//! A2A is an open protocol by Google Cloud for inter-agent communication.
//! This module implements the A2A client for AstrBot.
//!
//! Reference: https://google-a2a.github.io/
use crate::error::AstrBotError;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
// ============================================================================
// A2A Data Models
// ============================================================================
/// Agent Card - describes agent capabilities for service discovery
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AgentCard {
pub name: String,
pub description: Option<String>,
pub version: String,
pub provider: Option<AgentProvider>,
pub capabilities: AgentCapabilities,
pub skills: Vec<AgentSkill>,
#[serde(default)]
pub security_schemes: HashMap<String, SecurityScheme>,
pub url: String,
}
/// Provider information
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AgentProvider {
pub organization: String,
pub url: Option<String>,
}
/// Agent capabilities
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct AgentCapabilities {
#[serde(default)]
pub streaming: bool,
#[serde(default)]
pub push_notifications: bool,
#[serde(default)]
pub extensions: Vec<String>,
}
/// Agent skill
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AgentSkill {
pub id: String,
pub name: String,
#[serde(default)]
pub tags: Vec<String>,
#[serde(default)]
pub examples: Vec<SkillExample>,
}
/// Skill example
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SkillExample {
pub user: String,
pub agent: String,
}
/// Security scheme
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SecurityScheme {
#[serde(rename = "type")]
pub scheme_type: String,
#[serde(default)]
pub flows: Vec<String>,
#[serde(default)]
pub authorization_url: Option<String>,
}
/// Task status
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TaskStatus {
pub state: TaskState,
pub timestamp: String,
}
/// Task states
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum TaskState {
Submitted,
Working,
Completed,
Failed,
Canceled,
Rejected,
InputRequired,
AuthRequired,
}
/// Task representation
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Task {
pub id: String,
#[serde(default)]
pub context_id: Option<String>,
pub status: TaskStatus,
#[serde(default)]
pub artifacts: Vec<Artifact>,
#[serde(default)]
pub history: Vec<Message>,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
}
/// Artifact - immutable output from agent
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Artifact {
#[serde(default)]
pub artifact_id: Option<String>,
pub name: Option<String>,
pub parts: Vec<Part>,
}
/// Message part - atomic content unit
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(tag = "type")]
pub enum Part {
Text { text: String },
File { file: FilePart },
Data { data: serde_json::Value },
}
/// File part
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FilePart {
pub name: String,
pub uri: String,
}
/// Message between client and agent
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Message {
#[serde(default)]
pub message_id: Option<String>,
pub role: MessageRole,
pub parts: Vec<Part>,
#[serde(default)]
pub reference_task_ids: Vec<String>,
}
/// Message role
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum MessageRole {
User,
Agent,
}
/// A2A Response
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct A2AResponse {
#[serde(default)]
pub task: Option<Task>,
#[serde(default)]
pub message: Option<Message>,
}
// ============================================================================
// A2A Client
// ============================================================================
/// A2A Client for connecting to remote agents
#[derive(Debug)]
pub struct A2AClient {
/// Base URL of the A2A server
server_url: String,
/// Authentication header value
auth_header: Option<String>,
/// HTTP client
client: Option<reqwest::Client>,
/// Task history
tasks: HashMap<String, Task>,
}
impl A2AClient {
/// Create a new A2A client
pub fn new(server_url: &str) -> Self {
Self {
server_url: server_url.to_string(),
auth_header: None,
client: None,
tasks: HashMap::new(),
}
}
/// Set authentication
pub fn with_auth(mut self, auth_header: &str) -> Self {
self.auth_header = Some(auth_header.to_string());
self
}
/// Initialize the client
pub async fn connect(&mut self) -> Result<(), AstrBotError> {
self.client = Some(reqwest::Client::new());
tracing::info!("A2A client connected to {}", self.server_url);
Ok(())
}
/// Fetch agent card for service discovery
pub async fn get_agent_card(&self) -> Result<AgentCard, AstrBotError> {
let client = self
.client
.as_ref()
.ok_or_else(|| AstrBotError::NotConnected("A2A client not connected".into()))?;
let url = format!("{}/.well-known/agent-card.json", self.server_url);
let response = client.get(&url).send().await.map_err(|e| {
AstrBotError::ConnectionFailed(format!("Failed to fetch agent card: {}", e))
})?;
let agent_card: AgentCard = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(format!("Failed to parse agent card: {}", e)))?;
Ok(agent_card)
}
/// Send a message to an agent
pub async fn send_message(
&mut self,
message: Message,
task_id: Option<&str>,
session_id: Option<&str>,
) -> Result<A2AResponse, AstrBotError> {
let client = self
.client
.as_ref()
.ok_or_else(|| AstrBotError::NotConnected("A2A client not connected".into()))?;
let request = A2ARequest::new(
&uuid::Uuid::new_v4().to_string(),
"SendMessage",
serde_json::to_value(SendMessageParams {
message,
task_id: task_id.map(String::from),
session_id: session_id.map(String::from),
})
.map_err(|e| AstrBotError::Json(e))?,
);
let response = client
.post(&self.server_url)
.header("Content-Type", "application/json")
.header("A2A-Version", "1.0.0")
.json(&request)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(format!("SendMessage failed: {}", e)))?;
let result: serde_json::Value = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(format!("Failed to parse response: {}", e)))?;
// Parse into A2AResponse
let a2a_response: A2AResponse =
serde_json::from_value(result).map_err(|e| AstrBotError::Json(e))?;
Ok(a2a_response)
}
/// Send a streaming message
pub async fn send_streaming_message(
&mut self,
message: Message,
) -> Result<String, AstrBotError> {
let client = self
.client
.as_ref()
.ok_or_else(|| AstrBotError::NotConnected("A2A client not connected".into()))?;
let request = A2ARequest::new(
&uuid::Uuid::new_v4().to_string(),
"SendStreamingMessage",
serde_json::to_value(SendStreamingMessageParams { message })
.map_err(|e| AstrBotError::Json(e))?,
);
// For streaming, we return the task ID and let the caller handle SSE
let response = client
.post(&self.server_url)
.header("Content-Type", "application/json")
.header("A2A-Version", "1.0.0")
.json(&request)
.send()
.await
.map_err(|e| {
AstrBotError::ConnectionFailed(format!("SendStreamingMessage failed: {}", e))
})?;
// Parse streaming response to get task ID
let result: serde_json::Value = response.json().await.map_err(|e| {
AstrBotError::Protocol(format!("Failed to parse streaming response: {}", e))
})?;
let task_id = result
.get("task")
.and_then(|t| t.get("id"))
.and_then(|id| id.as_str())
.map(String::from)
.ok_or_else(|| AstrBotError::Protocol("No task ID in streaming response".into()))?;
Ok(task_id)
}
/// Get task status and result
pub async fn get_task(&mut self, task_id: &str) -> Result<Task, AstrBotError> {
let client = self
.client
.as_ref()
.ok_or_else(|| AstrBotError::NotConnected("A2A client not connected".into()))?;
let request = A2ARequest::new(
&uuid::Uuid::new_v4().to_string(),
"GetTask",
serde_json::json!({ "id": task_id }),
);
let response = client
.post(&self.server_url)
.header("Content-Type", "application/json")
.header("A2A-Version", "1.0.0")
.json(&request)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(format!("GetTask failed: {}", e)))?;
let result: serde_json::Value = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(format!("Failed to parse task: {}", e)))?;
let task: Task = serde_json::from_value(result).map_err(|e| AstrBotError::Json(e))?;
self.tasks.insert(task.id.clone(), task.clone());
Ok(task)
}
/// Cancel a task
pub async fn cancel_task(&mut self, task_id: &str) -> Result<Task, AstrBotError> {
let client = self
.client
.as_ref()
.ok_or_else(|| AstrBotError::NotConnected("A2A client not connected".into()))?;
let request = A2ARequest::new(
&uuid::Uuid::new_v4().to_string(),
"CancelTask",
serde_json::json!({ "id": task_id }),
);
let response = client
.post(&self.server_url)
.header("Content-Type", "application/json")
.header("A2A-Version", "1.0.0")
.json(&request)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(format!("CancelTask failed: {}", e)))?;
let result: serde_json::Value = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(format!("Failed to parse canceled task: {}", e)))?;
let task: Task = serde_json::from_value(result).map_err(|e| AstrBotError::Json(e))?;
Ok(task)
}
/// Check if agent is connected
pub fn is_connected(&self) -> bool {
self.client.is_some()
}
}
// ============================================================================
// A2A Request/Response Types
// ============================================================================
#[derive(Debug, Serialize, Deserialize)]
struct A2ARequest {
#[serde(rename = "jsonrpc")]
json_rpc: String,
id: String,
method: String,
params: serde_json::Value,
}
impl A2ARequest {
fn new(id: &str, method: &str, params: serde_json::Value) -> Self {
Self {
json_rpc: "2.0".to_string(),
id: id.to_string(),
method: method.to_string(),
params,
}
}
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct SendMessageParams {
message: Message,
#[serde(skip_serializing_if = "Option::is_none")]
task_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct SendStreamingMessageParams {
message: Message,
}
// ============================================================================
// Protocol Client Trait Implementation
// ============================================================================
#[async_trait]
impl crate::protocol::ProtocolClient for A2AClient {
fn name(&self) -> &'static str {
"a2a"
}
fn is_connected(&self) -> bool {
self.is_connected()
}
async fn connect(&mut self) -> Result<(), AstrBotError> {
A2AClient::connect(self).await
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.client = None;
self.tasks.clear();
tracing::info!("A2A client disconnected");
Ok(())
}
}

726
rust/src/abp.rs Normal file
View File

@@ -0,0 +1,726 @@
//! ABP (AstrBot Plugin) Protocol Implementation
//!
//! ABP is the plugin communication protocol for AstrBot, supporting:
//! - In-process and out-of-process loading modes
//! - Full plugin lifecycle management
//! - Tool calling, message handling, event subscriptions
//! - JSON-RPC 2.0 communication
//!
//! Reference: openspec/abp.md
use crate::error::AstrBotError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::UnixStream;
// ============================================================================
// ABP Data Models
// ============================================================================
/// Plugin loading mode
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PluginLoadMode {
/// In-process: direct function calls, zero serialization overhead
#[default]
InProcess,
/// Out-of-process: separate process, JSON-RPC over Unix Socket/HTTP
OutOfProcess,
}
impl std::fmt::Display for PluginLoadMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PluginLoadMode::InProcess => write!(f, "in_process"),
PluginLoadMode::OutOfProcess => write!(f, "out_of_process"),
}
}
}
/// Plugin transport type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PluginTransport {
/// Stdio transport (for subprocess)
#[default]
Stdio,
/// Unix Socket transport
UnixSocket,
/// HTTP transport
Http,
}
impl std::fmt::Display for PluginTransport {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
PluginTransport::Stdio => write!(f, "stdio"),
PluginTransport::UnixSocket => write!(f, "unix_socket"),
PluginTransport::Http => write!(f, "http"),
}
}
}
/// Plugin configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginConfig {
pub name: String,
#[serde(default)]
pub version: String,
#[serde(default)]
pub load_mode: PluginLoadMode,
#[serde(default)]
pub command: Option<String>,
#[serde(default)]
pub args: Vec<String>,
#[serde(default)]
pub env: HashMap<String, String>,
#[serde(default)]
pub transport: PluginTransport,
#[serde(default)]
pub url: Option<String>,
}
impl Default for PluginConfig {
fn default() -> Self {
Self {
name: String::new(),
version: "1.0.0".to_string(),
load_mode: PluginLoadMode::InProcess,
command: None,
args: Vec::new(),
env: HashMap::new(),
transport: PluginTransport::Stdio,
url: None,
}
}
}
/// Plugin capabilities
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginCapabilities {
#[serde(default)]
pub tools: bool,
#[serde(default)]
pub handlers: bool,
#[serde(default)]
pub events: bool,
#[serde(default)]
pub resources: bool,
}
/// Plugin metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PluginMetadata {
#[serde(default)]
pub display_name: Option<String>,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub author: Option<String>,
#[serde(default)]
pub homepage: Option<String>,
#[serde(default)]
pub support_platforms: Vec<String>,
#[serde(default)]
pub astrbot_version: Option<String>,
}
/// Initialize result from plugin
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InitializeResult {
pub protocol_version: String,
pub server_info: ServerInfo,
pub capabilities: PluginCapabilities,
#[serde(default)]
pub metadata: Option<PluginMetadata>,
}
/// Server info
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServerInfo {
pub name: String,
pub version: String,
}
/// Tool definition
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Tool {
pub name: String,
#[serde(default)]
pub description: String,
#[serde(default)]
pub input_schema: serde_json::Value,
}
/// Tool call result
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolResult {
#[serde(default)]
pub content: Vec<ToolContent>,
}
/// Tool content
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolContent {
#[serde(rename = "type")]
pub content_type: String,
pub text: Option<String>,
}
/// Message event
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MessageEvent {
pub message_id: String,
pub unified_msg_origin: String,
pub message_str: String,
pub sender: SenderInfo,
pub message_chain: Vec<MessageChainItem>,
}
/// Sender info
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SenderInfo {
pub user_id: String,
#[serde(default)]
pub nickname: String,
}
/// Message chain item
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct MessageChainItem {
#[serde(rename = "type")]
pub item_type: String,
#[serde(default)]
pub text: Option<String>,
}
/// Handle event result
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct HandleEventResult {
pub handled: bool,
#[serde(default)]
pub results: Vec<MessageChainItem>,
#[serde(default)]
pub stop_propagation: bool,
}
/// Event notification
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct EventNotification {
pub event_type: String,
#[serde(default)]
pub data: serde_json::Value,
}
// ============================================================================
// ABP Client
// ============================================================================
/// ABP Client for managing plugins
#[derive(Debug)]
pub struct AbpClient {
/// Connected state
connected: bool,
/// In-process plugins registry (name -> handler)
in_process_plugins: HashMap<String, InProcessPlugin>,
/// Out-of-process plugins registry (name -> OOP client)
out_of_process_plugins: HashMap<String, OutOfProcessPlugin>,
}
/// In-process plugin handler
#[derive(Debug, Clone)]
pub struct InProcessPlugin {
pub config: PluginConfig,
pub capabilities: PluginCapabilities,
pub metadata: Option<PluginMetadata>,
pub tools: Vec<Tool>,
}
/// Out-of-process plugin client
#[derive(Debug)]
pub struct OutOfProcessPlugin {
pub config: PluginConfig,
pub capabilities: PluginCapabilities,
pub metadata: Option<PluginMetadata>,
pub tools: Vec<Tool>,
pub socket_path: Option<PathBuf>,
pub http_url: Option<String>,
}
impl AbpClient {
/// Create a new ABP client
pub fn new() -> Self {
Self {
connected: false,
in_process_plugins: HashMap::new(),
out_of_process_plugins: HashMap::new(),
}
}
/// Connect the ABP client
pub async fn connect(&mut self) -> Result<(), AstrBotError> {
self.connected = true;
tracing::info!(
"ABP client connected (in_process: {}, out_of_process: {})",
self.in_process_plugins.len(),
self.out_of_process_plugins.len()
);
Ok(())
}
/// Disconnect the ABP client
pub async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.connected = false;
self.in_process_plugins.clear();
self.out_of_process_plugins.clear();
tracing::info!("ABP client disconnected");
Ok(())
}
/// Check if connected
pub fn is_connected(&self) -> bool {
self.connected
}
/// Set connected state (sync version for orchestrator)
pub fn set_connected(&mut self, connected: bool) {
self.connected = connected;
}
/// Set connected state without consuming self
pub fn set_connected_ref(&mut self, connected: bool) {
self.connected = connected;
}
/// Register an in-process plugin
pub fn register_in_process_plugin(&mut self, config: PluginConfig) {
let plugin = InProcessPlugin {
config: config.clone(),
capabilities: PluginCapabilities::default(),
metadata: None,
tools: Vec::new(),
};
tracing::debug!("Registered in-process plugin: {}", config.name);
self.in_process_plugins.insert(config.name, plugin);
}
/// Register an out-of-process plugin
pub fn register_out_of_process_plugin(&mut self, config: PluginConfig) {
let socket_path = if config.transport == PluginTransport::UnixSocket {
Some(PathBuf::from(format!(
"/tmp/astrbot_plugin_{}.sock",
config.name
)))
} else {
None
};
let plugin = OutOfProcessPlugin {
config: config.clone(),
capabilities: PluginCapabilities::default(),
metadata: None,
tools: Vec::new(),
socket_path,
http_url: config.url.clone(),
};
tracing::debug!("Registered out-of-process plugin: {}", config.name);
self.out_of_process_plugins.insert(config.name, plugin);
}
/// Unregister a plugin
pub fn unregister_plugin(&mut self, name: &str) {
self.in_process_plugins.remove(name);
self.out_of_process_plugins.remove(name);
tracing::debug!("Unregistered plugin: {}", name);
}
/// List all registered plugin names
pub fn list_plugins(&self) -> Vec<String> {
let mut names: Vec<String> = self.in_process_plugins.keys().cloned().collect();
names.extend(self.out_of_process_plugins.keys().cloned());
names.sort();
names.dedup();
names
}
/// Get plugin info
pub fn get_plugin_info(&self, name: &str) -> Option<PluginInfo> {
if let Some(plugin) = self.in_process_plugins.get(name) {
Some(PluginInfo {
name: plugin.config.name.clone(),
version: plugin.config.version.clone(),
load_mode: plugin.config.load_mode,
capabilities: plugin.capabilities.clone(),
metadata: plugin.metadata.clone(),
tools_count: plugin.tools.len(),
})
} else if let Some(plugin) = self.out_of_process_plugins.get(name) {
Some(PluginInfo {
name: plugin.config.name.clone(),
version: plugin.config.version.clone(),
load_mode: plugin.config.load_mode,
capabilities: plugin.capabilities.clone(),
metadata: plugin.metadata.clone(),
tools_count: plugin.tools.len(),
})
} else {
None
}
}
/// Call a plugin tool (out-of-process)
pub async fn call_tool(
&mut self,
plugin_name: &str,
tool_name: &str,
arguments: serde_json::Value,
) -> Result<ToolResult, AstrBotError> {
// Extract transport and connection info first to avoid borrow conflict
let transport = {
let plugin = self
.out_of_process_plugins
.get(plugin_name)
.ok_or_else(|| {
AstrBotError::NotFound(format!("Plugin '{}' not found", plugin_name))
})?;
plugin.config.transport
};
match transport {
PluginTransport::UnixSocket => {
let socket_path = {
let plugin = self
.out_of_process_plugins
.get(plugin_name)
.ok_or_else(|| {
AstrBotError::NotFound(format!("Plugin '{}' not found", plugin_name))
})?;
plugin.socket_path.clone()
};
self.call_tool_unix_socket(&socket_path, tool_name, arguments)
.await
}
PluginTransport::Http => {
let http_url = {
let plugin = self
.out_of_process_plugins
.get(plugin_name)
.ok_or_else(|| {
AstrBotError::NotFound(format!("Plugin '{}' not found", plugin_name))
})?;
plugin.http_url.clone()
};
self.call_tool_http(&http_url, tool_name, arguments).await
}
PluginTransport::Stdio => Err(AstrBotError::Protocol(
"Stdio transport not implemented for tool calls".into(),
)),
}
}
/// Call tool via Unix Socket
async fn call_tool_unix_socket(
&self,
socket_path: &Option<PathBuf>,
tool_name: &str,
arguments: serde_json::Value,
) -> Result<ToolResult, AstrBotError> {
let socket_path = socket_path
.as_ref()
.ok_or_else(|| AstrBotError::InvalidState("No socket path configured".into()))?;
let mut stream = UnixStream::connect(socket_path).await.map_err(|e| {
AstrBotError::ConnectionFailed(format!("Unix socket connect failed: {}", e))
})?;
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": uuid::Uuid::new_v4().to_string(),
"method": "tools/call",
"params": {
"name": tool_name,
"arguments": arguments
}
});
let content = serde_json::to_string(&request).map_err(|e| AstrBotError::Json(e))?;
let header = format!("Content-Length: {}\r\n\r\n", content.len());
stream.write_all(header.as_bytes()).await?;
stream.write_all(content.as_bytes()).await?;
// Read response
let mut buffer = Vec::new();
stream.read_to_end(&mut buffer).await?;
let response: serde_json::Value =
serde_json::from_slice(&buffer).map_err(|e| AstrBotError::Json(e))?;
if let Some(error) = response.get("error") {
return Err(AstrBotError::Protocol(error.to_string()));
}
let result: ToolResult =
serde_json::from_value(response.get("result").cloned().unwrap_or_default())
.map_err(|e| AstrBotError::Json(e))?;
Ok(result)
}
/// Call tool via HTTP
async fn call_tool_http(
&self,
http_url: &Option<String>,
tool_name: &str,
arguments: serde_json::Value,
) -> Result<ToolResult, AstrBotError> {
let url = http_url
.as_ref()
.ok_or_else(|| AstrBotError::InvalidState("No HTTP URL configured".into()))?;
let client = reqwest::Client::new();
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": uuid::Uuid::new_v4().to_string(),
"method": "tools/call",
"params": {
"name": tool_name,
"arguments": arguments
}
});
let response = client
.post(url)
.json(&request)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(e.to_string()))?;
let result: ToolResult = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(e.to_string()))?;
Ok(result)
}
/// Handle an event with a plugin
pub async fn handle_event(
&mut self,
plugin_name: &str,
event_type: &str,
event: serde_json::Value,
) -> Result<HandleEventResult, AstrBotError> {
// Extract transport first to avoid borrow conflict
let transport = {
let plugin = self
.out_of_process_plugins
.get(plugin_name)
.ok_or_else(|| {
AstrBotError::NotFound(format!("Plugin '{}' not found", plugin_name))
})?;
plugin.config.transport
};
match transport {
PluginTransport::UnixSocket => {
let socket_path = {
let plugin = self
.out_of_process_plugins
.get(plugin_name)
.ok_or_else(|| {
AstrBotError::NotFound(format!("Plugin '{}' not found", plugin_name))
})?;
plugin.socket_path.clone()
};
self.handle_event_unix_socket(&socket_path, event_type, event)
.await
}
PluginTransport::Http => {
let http_url = {
let plugin = self
.out_of_process_plugins
.get(plugin_name)
.ok_or_else(|| {
AstrBotError::NotFound(format!("Plugin '{}' not found", plugin_name))
})?;
plugin.http_url.clone()
};
self.handle_event_http(&http_url, event_type, event).await
}
PluginTransport::Stdio => Err(AstrBotError::Protocol(
"Stdio transport not implemented for events".into(),
)),
}
}
/// Handle event via Unix Socket
async fn handle_event_unix_socket(
&self,
socket_path: &Option<PathBuf>,
event_type: &str,
event: serde_json::Value,
) -> Result<HandleEventResult, AstrBotError> {
let socket_path = socket_path
.as_ref()
.ok_or_else(|| AstrBotError::InvalidState("No socket path configured".into()))?;
let mut stream = UnixStream::connect(socket_path).await.map_err(|e| {
AstrBotError::ConnectionFailed(format!("Unix socket connect failed: {}", e))
})?;
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": uuid::Uuid::new_v4().to_string(),
"method": "plugin.handle_event",
"params": {
"event_type": event_type,
"event": event
}
});
let content = serde_json::to_string(&request).map_err(|e| AstrBotError::Json(e))?;
let header = format!("Content-Length: {}\r\n\r\n", content.len());
stream.write_all(header.as_bytes()).await?;
stream.write_all(content.as_bytes()).await?;
// Read response
let mut buffer = Vec::new();
stream.read_to_end(&mut buffer).await?;
let response: serde_json::Value =
serde_json::from_slice(&buffer).map_err(|e| AstrBotError::Json(e))?;
if let Some(error) = response.get("error") {
return Err(AstrBotError::Protocol(error.to_string()));
}
let result: HandleEventResult =
serde_json::from_value(response.get("result").cloned().unwrap_or_default())
.map_err(|e| AstrBotError::Json(e))?;
Ok(result)
}
/// Handle event via HTTP
async fn handle_event_http(
&self,
http_url: &Option<String>,
event_type: &str,
event: serde_json::Value,
) -> Result<HandleEventResult, AstrBotError> {
let url = http_url
.as_ref()
.ok_or_else(|| AstrBotError::InvalidState("No HTTP URL configured".into()))?;
let client = reqwest::Client::new();
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": uuid::Uuid::new_v4().to_string(),
"method": "plugin.handle_event",
"params": {
"event_type": event_type,
"event": event
}
});
let response = client
.post(url)
.json(&request)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(e.to_string()))?;
let result: HandleEventResult = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(e.to_string()))?;
Ok(result)
}
/// Get plugin tools
pub fn get_plugin_tools(&self, name: &str) -> Option<Vec<Tool>> {
if let Some(plugin) = self.in_process_plugins.get(name) {
Some(plugin.tools.clone())
} else if let Some(plugin) = self.out_of_process_plugins.get(name) {
Some(plugin.tools.clone())
} else {
None
}
}
/// Health check for a plugin
pub fn health_check(&self, name: &str) -> bool {
if self.in_process_plugins.contains_key(name) {
true
} else if let Some(plugin) = self.out_of_process_plugins.get(name) {
if let Some(socket_path) = &plugin.socket_path {
socket_path.exists()
} else if plugin.http_url.is_some() {
true // HTTP plugins are assumed healthy if URL is configured
} else {
false
}
} else {
false
}
}
}
impl Default for AbpClient {
fn default() -> Self {
Self::new()
}
}
/// Plugin info for listing
#[derive(Debug, Clone)]
pub struct PluginInfo {
pub name: String,
pub version: String,
pub load_mode: PluginLoadMode,
pub capabilities: PluginCapabilities,
pub metadata: Option<PluginMetadata>,
pub tools_count: usize,
}
// ============================================================================
// Protocol Client Trait Implementation
// ============================================================================
use async_trait::async_trait;
#[async_trait]
impl crate::protocol::ProtocolClient for AbpClient {
fn name(&self) -> &'static str {
"abp"
}
fn is_connected(&self) -> bool {
self.is_connected()
}
async fn connect(&mut self) -> Result<(), AstrBotError> {
self.connect().await
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.disconnect().await
}
}

View File

@@ -1,7 +1,7 @@
//! AstrBot Core CLI
use anyhow::Result;
use clap::{Parser, Subcommand};
use clap::{CommandFactory, Parser, Subcommand};
#[derive(Parser)]
#[command(name = "astrbot-rs")]
@@ -13,7 +13,7 @@ pub struct Cli {
#[derive(Subcommand)]
pub enum Commands {
/// Start the AstrBot core runtime
/// Start the AstrBot Core runtime
Start {
/// Host to bind to
#[arg(long, default_value = "127.0.0.1")]
@@ -34,7 +34,64 @@ pub fn cli() -> Result<()> {
}
pub fn cli_with_args(args: &[String]) -> Result<()> {
let cli = Cli::try_parse_from(args).map_err(|e| anyhow::anyhow!("{e}"))?;
// Handle --help and --version explicitly
if args.iter().any(|a| a == "--help" || a == "-h") {
let mut cmd = Cli::command();
cmd.print_help()?;
println!();
return Ok(());
}
if args.iter().any(|a| a == "--version" || a == "-V") {
println!("astrbot-rs {}", env!("CARGO_PKG_VERSION"));
return Ok(());
}
// Handle empty args - show help
if args.is_empty() {
let mut cmd = Cli::command();
cmd.print_help()?;
println!();
return Ok(());
}
// When called from Python with sys.argv, clap treats the first element as bin_name.
// We need to handle two cases:
// 1. Single element that's a known subcommand (e.g., ['stats']) - prepend bin_name
// 2. Multiple elements where first is program name (e.g., ['/path/astrbot-rs', 'stats']) - prepend bin_name
let known_subcommands = ["start", "stats", "health", "help"];
let first_is_subcommand = args.len() == 1 && known_subcommands.contains(&args[0].as_str());
let first_is_program_name = !args[0].starts_with('-');
let parse_args: Vec<&str> = if first_is_subcommand {
// Case 1: ['stats'] -> prepend bin_name
vec!["astrbot-rs", args[0].as_str()]
} else if first_is_program_name {
// Case 2: ['/path/astrbot-rs', 'stats', ...] -> strip first, prepend bin_name
if args.len() > 1 {
let mut prefixed = vec!["astrbot-rs"];
prefixed.extend(args[1..].iter().map(|s| s.as_str()));
prefixed
} else {
// Only program name, no subcommand - show help
let mut cmd = Cli::command();
cmd.print_help()?;
println!();
return Ok(());
}
} else {
// Normal case: args start with flags
args.iter().map(|s| s.as_str()).collect()
};
// If all args consumed, show help
if parse_args.is_empty() {
let mut cmd = Cli::command();
cmd.print_help()?;
println!();
return Ok(());
}
let cli = Cli::try_parse_from(parse_args).map_err(|e| anyhow::anyhow!("{e}"))?;
match cli.command {
Commands::Start { host, port } => {

File diff suppressed because it is too large Load Diff

View File

@@ -19,6 +19,9 @@ pub enum AstrBotError {
#[error("Invalid state: {0}")]
InvalidState(String),
#[error("Not found: {0}")]
NotFound(String),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),

View File

@@ -1,15 +1,14 @@
//! AstrBot Core - High-performance core runtime in Rust
//! AstrBot Core - High-performance runtime in Rust
//!
//! This crate provides the core runtime components for AstrBot,
//! exposing Python bindings via pyo3.
//! This crate provides the core runtime for AstrBot, with Rust as the core
//! and Python modules as plugins via PyO3.
// RULES:
// - NO unsafe blocks allowed
// - NO .unwrap() - use ? or expect with message
// - All errors must be handled properly
// - NO .unwrap() without message - use ? or expect()
// - All errors must be handled properly via Result
#![deny(clippy::all)]
#![deny(clippy::pedantic)]
#![deny(unsafe_code)]
#![allow(clippy::module_name_repetitions)]
#![allow(clippy::too_many_lines)]
@@ -21,23 +20,32 @@
#![allow(clippy::missing_errors_doc)]
#![allow(clippy::missing_panics_doc)]
#![allow(clippy::doc_markdown)]
#![allow(clippy::unwrap_used)]
#![allow(clippy::cast_precision_loss)]
#![allow(clippy::return_self_not_must_use)]
#![allow(clippy::must_use_candidate)]
pub mod a2a;
pub mod abp;
pub mod cli;
pub mod error;
pub mod orchestrator;
pub mod message;
pub mod stats;
pub mod protocol;
pub mod config;
pub mod error;
pub mod message;
pub mod orchestrator;
pub mod protocol;
pub mod server;
pub mod stats;
#[cfg(feature = "python")]
pub mod python;
pub use a2a::{AgentCard, Task, TaskState};
pub use abp::{PluginCapabilities, PluginConfig, PluginLoadMode};
pub use error::AstrBotError;
pub use orchestrator::Orchestrator;
pub use message::Message;
pub use orchestrator::Orchestrator;
pub use protocol::ProtocolStatus;
pub use server::{ApiResponse, HttpServer, WsManager};
pub use stats::RuntimeStats;
// Re-export CLI for Python bindings
#[cfg(feature = "python")]
pub use cli::cli_with_args;

View File

@@ -1,33 +1,43 @@
//! Core orchestrator for AstrBot runtime
//!
//! Manages lifecycle of all protocol clients and stars (plugins).
use crate::abp::{AbpClient, PluginConfig, PluginLoadMode};
use crate::error::AstrBotError;
use crate::protocol::{AcpClient, LspClient, McpClient, ProtocolClient};
use crate::stats::RuntimeStats;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use tokio::sync::broadcast;
use tokio::time::{interval, Duration};
use tokio::time::{Duration, interval};
use tracing::{debug, info, warn};
#[derive(Debug, Clone)]
pub struct ProtocolStatus {
pub connected: bool,
pub name: String,
}
// ============================================================================
// Orchestrator
// ============================================================================
impl Default for ProtocolStatus {
fn default() -> Self {
Self {
connected: false,
name: String::new(),
}
}
}
/// Main orchestrator coordinating all protocol clients
/// Main orchestrator coordinating all protocol clients and stars
pub struct Orchestrator {
/// Running state
running: RwLock<bool>,
/// Shutdown signal sender
shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
stars: RwLock<HashMap<String, String>>,
stats: RwLock<RuntimeStats>,
/// Protocol clients
lsp: RwLock<LspClient>,
mcp: RwLock<McpClient>,
acp: RwLock<AcpClient>,
abp: RwLock<AbpClient>,
/// Star registry
stars: RwLock<HashMap<String, StarRegistration>>,
/// Runtime statistics
stats: RuntimeStats,
}
/// Star registration entry
#[derive(Debug, Clone)]
pub struct StarRegistration {
pub name: String,
pub handler: String,
}
impl Default for Orchestrator {
@@ -37,52 +47,108 @@ impl Default for Orchestrator {
}
impl Orchestrator {
/// Create a new Orchestrator instance
#[must_use]
pub fn new() -> Self {
Self {
running: RwLock::new(false),
shutdown_tx: Arc::new(RwLock::new(None)),
lsp: RwLock::new(LspClient::new()),
mcp: RwLock::new(McpClient::new()),
acp: RwLock::new(AcpClient::new()),
abp: RwLock::new(AbpClient::new()),
stars: RwLock::new(HashMap::new()),
stats: RwLock::new(RuntimeStats::default()),
stats: RuntimeStats::new(),
}
}
/// Start the orchestrator
pub fn start(&self) -> Result<(), AstrBotError> {
/// Start the orchestrator and all protocol clients (sync version)
pub fn start_sync(&self) -> Result<(), AstrBotError> {
{
let mut running = self.running.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
let mut running = self
.running
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
if *running {
return Err(AstrBotError::InvalidState(
"Orchestrator already started".into(),
));
}
*running = true;
}
let (tx, _rx) = broadcast::channel(1);
{
let mut shutdown_tx = self.shutdown_tx.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
let mut shutdown_tx = self
.shutdown_tx
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
*shutdown_tx = Some(tx);
}
tracing::info!("Orchestrator started");
// Connect all protocol clients (sync)
self.connect_protocols_sync()?;
info!("Orchestrator started");
Ok(())
}
/// Bootstrap all protocol clients
pub async fn bootstrap(&self) -> Result<(), AstrBotError> {
self.start()?;
tracing::info!("Protocol clients would be started here");
tokio::time::sleep(Duration::from_millis(500)).await;
/// Connect all protocol clients (sync version for Python binding)
fn connect_protocols_sync(&self) -> Result<(), AstrBotError> {
// Connect LSP
{
let mut lsp = self
.lsp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
// For now, just mark as connected (actual connection is async)
lsp.set_connected(true);
}
// Connect MCP
{
let mut mcp = self
.mcp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
mcp.set_connected(true);
}
// Connect ACP
{
let mut acp = self
.acp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
acp.set_connected(true);
}
// Connect ABP
{
let mut abp = self
.abp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
abp.set_connected(true);
}
Ok(())
}
/// Start the orchestrator and all protocol clients (async)
pub async fn start(&self) -> Result<(), AstrBotError> {
self.start_sync()
}
/// Main event loop
pub async fn run_loop(&self) -> Result<(), AstrBotError> {
if !self.is_running() {
return Err(AstrBotError::InvalidState("Orchestrator not started".into()));
return Err(AstrBotError::InvalidState(
"Orchestrator not started".into(),
));
}
tracing::info!("Orchestrator event loop started");
info!("Orchestrator event loop started");
let mut tick_interval = interval(Duration::from_secs(5));
loop {
@@ -91,7 +157,7 @@ impl Orchestrator {
self.periodic_health_check();
}
_ = self.wait_for_shutdown() => {
tracing::info!("Orchestrator shutdown signal received");
info!("Orchestrator shutdown signal received");
break;
}
}
@@ -101,10 +167,11 @@ impl Orchestrator {
}
}
tracing::info!("Orchestrator event loop stopped");
info!("Orchestrator event loop stopped");
Ok(())
}
/// Wait for shutdown signal
async fn wait_for_shutdown(&self) {
let tx_guard = self.shutdown_tx.read().ok();
let tx = tx_guard.as_ref().and_then(|t| t.as_ref());
@@ -117,16 +184,36 @@ impl Orchestrator {
}
}
/// Periodic health check for all protocol clients
fn periodic_health_check(&self) {
tracing::debug!("Orchestrator health check");
debug!("Running periodic health check");
if let Ok(lsp) = self.lsp.read() {
if !lsp.is_connected() {
warn!("LSP client disconnected");
}
}
if let Ok(mcp) = self.mcp.read() {
if !mcp.is_connected() {
warn!("MCP client disconnected");
}
}
if let Ok(acp) = self.acp.read() {
if !acp.is_connected() {
warn!("ACP client disconnected");
}
}
}
/// Stop the orchestrator
pub fn stop(&self) -> Result<(), AstrBotError> {
/// Stop the orchestrator (sync version)
pub fn stop_sync(&self) -> Result<(), AstrBotError> {
{
let mut running = self.running.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
let mut running = self
.running
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
*running = false;
}
@@ -136,31 +223,121 @@ impl Orchestrator {
}
}
tracing::info!("Orchestrator stopped");
self.shutdown_protocols_sync()?;
info!("Orchestrator stopped");
Ok(())
}
/// Shutdown all protocol clients (sync version)
fn shutdown_protocols_sync(&self) -> Result<(), AstrBotError> {
{
let mut lsp = self
.lsp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
lsp.set_connected(false);
}
{
let mut mcp = self
.mcp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
mcp.set_connected(false);
}
{
let mut acp = self
.acp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
acp.set_connected(false);
}
{
let mut abp = self
.abp
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
abp.set_connected(false);
}
Ok(())
}
/// Stop the orchestrator (async)
pub async fn stop(&self) -> Result<(), AstrBotError> {
self.stop_sync()
}
/// Check if orchestrator is running
#[must_use]
pub fn is_running(&self) -> bool {
self.running.read().map(|r| *r).unwrap_or(false)
}
pub fn register_star(&self, name: &str, _handler: &str) -> Result<(), AstrBotError> {
let mut stars = self.stars.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
stars.insert(name.to_string(), name.to_string());
/// Register a star (plugin)
pub fn register_star(&self, name: &str, handler: &str) -> Result<(), AstrBotError> {
let mut stars = self
.stars
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
let registration = StarRegistration {
name: name.to_string(),
handler: handler.to_string(),
};
stars.insert(name.to_string(), registration);
// 根据 handler 判断加载模式:包含 "/" 视为 Unix Socket 路径,否则为模块名
if handler.starts_with('/') || handler.contains(".sock") {
// 跨进程加载
if let Ok(mut abp) = self.abp.write() {
let config = PluginConfig {
name: name.to_string(),
version: "1.0.0".to_string(),
load_mode: PluginLoadMode::OutOfProcess,
command: Some(handler.to_string()),
..Default::default()
};
abp.register_out_of_process_plugin(config);
}
} else {
// 进程内加载
if let Ok(mut abp) = self.abp.write() {
let config = PluginConfig {
name: name.to_string(),
version: "1.0.0".to_string(),
load_mode: PluginLoadMode::InProcess,
..Default::default()
};
abp.register_in_process_plugin(config);
}
}
info!("Star '{}' registered", name);
Ok(())
}
/// Unregister a star (plugin)
pub fn unregister_star(&self, name: &str) -> Result<(), AstrBotError> {
let mut stars = self.stars.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
let mut stars = self
.stars
.write()
.map_err(|_| AstrBotError::InvalidState("Failed to acquire write lock".into()))?;
stars.remove(name);
if let Ok(mut abp) = self.abp.write() {
abp.unregister_plugin(name);
}
info!("Star '{}' unregistered", name);
Ok(())
}
/// List all registered stars
#[must_use]
pub fn list_stars(&self) -> Vec<String> {
self.stars
@@ -169,26 +346,196 @@ impl Orchestrator {
.unwrap_or_default()
}
/// Record a message activity
pub fn record_activity(&self) {
if let Ok(stats) = self.stats.write() {
stats.record_message();
self.stats.record_message();
}
/// Get runtime statistics
#[must_use]
pub fn stats(&self) -> RuntimeStats {
self.stats.clone()
}
/// Set protocol connection status
pub fn set_protocol_connected(
&self,
protocol: &str,
connected: bool,
) -> Result<(), AstrBotError> {
match protocol {
"lsp" => {
let mut lsp = self.lsp.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
lsp.set_connected(connected);
}
"mcp" => {
let mut mcp = self.mcp.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
mcp.set_connected(connected);
}
"acp" => {
let mut acp = self.acp.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
acp.set_connected(connected);
}
"abp" => {
let mut abp = self.abp.write().map_err(|_| {
AstrBotError::InvalidState("Failed to acquire write lock".into())
})?;
abp.set_connected(connected);
}
_ => {
return Err(AstrBotError::InvalidState(format!(
"Unknown protocol: {protocol}"
)));
}
}
Ok(())
}
/// Get protocol status
#[must_use]
pub fn get_protocol_status(&self, protocol: &str) -> Option<crate::protocol::ProtocolStatus> {
match protocol {
"lsp" => self
.lsp
.read()
.ok()
.map(|lsp| crate::protocol::ProtocolStatus {
connected: lsp.is_connected(),
name: "lsp".to_string(),
}),
"mcp" => self
.mcp
.read()
.ok()
.map(|mcp| crate::protocol::ProtocolStatus {
connected: mcp.is_connected(),
name: "mcp".to_string(),
}),
"acp" => self
.acp
.read()
.ok()
.map(|acp| crate::protocol::ProtocolStatus {
connected: acp.is_connected(),
name: "acp".to_string(),
}),
"abp" => self
.abp
.read()
.ok()
.map(|abp| crate::protocol::ProtocolStatus {
connected: abp.is_connected(),
name: "abp".to_string(),
}),
_ => None,
}
}
}
// ============================================================================
// Python bindings via PyO3 (sync only)
// ============================================================================
#[cfg(feature = "python")]
mod python {
use super::*;
use pyo3::prelude::*;
use pyo3::types::PyDict;
/// Python wrapper for Orchestrator
#[pyclass]
pub struct PyOrchestrator {
inner: Orchestrator,
}
#[pymethods]
impl PyOrchestrator {
#[new]
pub fn new() -> Self {
Self {
inner: Orchestrator::new(),
}
}
/// Start the orchestrator (sync, callable from Python)
pub fn start(&self) -> PyResult<()> {
self.inner
.start_sync()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
/// Stop the orchestrator (sync, callable from Python)
pub fn stop(&self) -> PyResult<()> {
self.inner
.stop_sync()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
pub fn is_running(&self) -> bool {
self.inner.is_running()
}
pub fn register_star(&self, name: &str, handler: &str) -> PyResult<()> {
self.inner
.register_star(name, handler)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
pub fn unregister_star(&self, name: &str) -> PyResult<()> {
self.inner
.unregister_star(name)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
pub fn list_stars(&self) -> Vec<String> {
self.inner.list_stars()
}
pub fn record_activity(&self) {
self.inner.record_activity();
}
pub fn get_stats(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
let stats = self.inner.stats();
let dict = PyDict::new(py);
dict.set_item("message_count", stats.message_count())?;
dict.set_item("uptime_seconds", stats.uptime_seconds())?;
dict.set_item("last_activity_time", stats.last_activity_time())?;
Ok(dict.into())
}
pub fn set_protocol_connected(&self, protocol: &str, connected: bool) -> PyResult<()> {
self.inner
.set_protocol_connected(protocol, connected)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
pub fn get_protocol_status(
&self,
protocol: &str,
py: Python<'_>,
) -> PyResult<Option<Py<PyDict>>> {
Ok(self.inner.get_protocol_status(protocol).map(|status| {
let dict = PyDict::new(py);
dict.set_item("connected", status.connected).unwrap();
dict.set_item("name", &status.name).unwrap();
dict.into()
}))
}
}
#[must_use]
pub fn stats(&self) -> RuntimeStats {
self.stats.read().map(|s| s.clone()).unwrap_or_default()
}
#[must_use]
pub fn get_protocol_status(&self, _protocol: &str) -> Option<ProtocolStatus> {
Some(ProtocolStatus {
connected: true,
name: _protocol.to_string(),
})
}
pub fn set_protocol_connected(&self, _protocol: &str, _connected: bool) -> Result<(), AstrBotError> {
Ok(())
impl Default for PyOrchestrator {
fn default() -> Self {
Self::new()
}
}
}
#[cfg(feature = "python")]
pub use python::PyOrchestrator;

View File

@@ -1,119 +1,67 @@
//! Protocol client implementations for ACP and ABP
//! Protocol client implementations for LSP, MCP, ACP, and ABP
use crate::error::AstrBotError;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::net::UnixStream;
use tokio::process::{Child, Command};
/// Protocol client trait
#[async_trait]
pub trait ProtocolClient: Send + Sync {
async fn connect(&mut self) -> Result<(), AstrBotError>;
async fn disconnect(&mut self) -> Result<(), AstrBotError>;
fn is_connected(&self) -> bool;
fn name(&self) -> &str;
// ============================================================================
// Protocol Status
// ============================================================================
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProtocolStatus {
pub connected: bool,
pub name: String,
}
impl ProtocolStatus {
pub fn new(name: &str, connected: bool) -> Self {
Self {
connected,
name: name.to_string(),
}
}
}
// ============================================================================
// LSP Client
// Protocol Client Trait
// ============================================================================
#[async_trait]
pub trait ProtocolClient: Send + Sync {
fn name(&self) -> &str;
fn is_connected(&self) -> bool;
async fn connect(&mut self) -> Result<(), AstrBotError>;
async fn disconnect(&mut self) -> Result<(), AstrBotError>;
}
// ============================================================================
// LSP Client - Language Server Protocol
// ============================================================================
#[derive(Default)]
pub struct LspClient {
connected: bool,
child: Option<Child>,
stdin: Option<tokio::io::BufWriter<tokio::process::ChildStdin>>,
stdout: Option<tokio::io::BufReader<tokio::process::ChildStdout>>,
pending_requests: HashMap<i64, tokio::sync::oneshot::Sender<String>>,
request_id: i64,
}
impl LspClient {
#[must_use]
pub fn new() -> Self {
Self { connected: false }
}
pub fn set_connected(&mut self, connected: bool) {
self.connected = connected;
}
}
#[async_trait]
impl ProtocolClient for LspClient {
async fn connect(&mut self) -> Result<(), AstrBotError> {
tracing::debug!("LSP client connecting");
self.connected = true;
Ok(())
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
tracing::debug!("LSP client disconnecting");
self.connected = false;
Ok(())
}
fn is_connected(&self) -> bool {
self.connected
}
fn name(&self) -> &'static str {
"lsp-client"
}
}
// ============================================================================
// MCP Client
// ============================================================================
#[derive(Default)]
pub struct McpClient {
connected: bool,
}
impl McpClient {
#[must_use]
pub fn new() -> Self {
Self { connected: false }
}
pub fn set_connected(&mut self, connected: bool) {
self.connected = connected;
}
}
#[async_trait]
impl ProtocolClient for McpClient {
async fn connect(&mut self) -> Result<(), AstrBotError> {
tracing::debug!("MCP client connecting");
self.connected = true;
Ok(())
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
tracing::debug!("MCP client disconnecting");
self.connected = false;
Ok(())
}
fn is_connected(&self) -> bool {
self.connected
}
fn name(&self) -> &'static str {
"mcp-client"
}
}
// ============================================================================
// ACP Client - AstrBot Communication Protocol
// ============================================================================
#[derive(Debug, Clone)]
pub struct AcpClient {
connected: bool,
server_url: Option<String>,
}
impl AcpClient {
pub fn new() -> Self {
Self {
connected: false,
server_url: None,
child: None,
stdin: None,
stdout: None,
pending_requests: HashMap::new(),
request_id: 0,
}
}
@@ -121,17 +69,512 @@ impl AcpClient {
self.connected = connected;
}
pub async fn connect_to_server(&mut self, host: &str, port: u16) -> Result<(), AstrBotError> {
self.server_url = Some(format!("{}:{}", host, port));
/// Connect to an LSP server subprocess
pub async fn connect_to_server(
&mut self,
command: Vec<String>,
workspace_uri: &str,
) -> Result<(), AstrBotError> {
if command.is_empty() {
return Err(AstrBotError::InvalidState("LSP command is empty".into()));
}
let mut cmd = Command::new(&command[0]);
if command.len() > 1 {
cmd.args(&command[1..]);
}
cmd.stdin(std::process::Stdio::piped());
cmd.stdout(std::process::Stdio::piped());
cmd.stderr(std::process::Stdio::piped());
let mut child = cmd.spawn().map_err(|e| {
AstrBotError::ConnectionFailed(format!("Failed to spawn LSP server: {e}"))
})?;
let stdin = child
.stdin
.take()
.ok_or_else(|| AstrBotError::ConnectionFailed("Failed to capture stdin".into()))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| AstrBotError::ConnectionFailed("Failed to capture stdout".into()))?;
self.stdin = Some(tokio::io::BufWriter::new(stdin));
self.stdout = Some(tokio::io::BufReader::new(stdout));
self.child = Some(child);
self.connected = true;
tracing::debug!("ACP client connecting to {}:{}", host, port);
// Send initialize request
let initialize_params = serde_json::json!({
"processId": std::process::id(),
"rootUri": workspace_uri,
"capabilities": {}
});
self.send_request("initialize", Some(initialize_params))
.await?;
// Send initialized notification
self.send_notification("initialized", None).await?;
tracing::info!("LSP client connected to server");
Ok(())
}
pub async fn connect_to_unix_socket(&mut self, socket_path: &str) -> Result<(), AstrBotError> {
self.server_url = Some(format!("unix://{}", socket_path));
/// Send a request and wait for response
pub async fn send_request(
&mut self,
method: &str,
params: Option<serde_json::Value>,
) -> Result<String, AstrBotError> {
let (tx, rx) = tokio::sync::oneshot::channel();
let id = self.request_id;
self.request_id += 1;
self.pending_requests.insert(id, tx);
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"method": method,
"params": params.unwrap_or(serde_json::Value::Null)
});
self.send_json(&request).await?;
rx.await
.map_err(|_| AstrBotError::Timeout("LSP request timed out".into()))
}
/// Send a notification (no response expected)
pub async fn send_notification(
&mut self,
method: &str,
params: Option<serde_json::Value>,
) -> Result<(), AstrBotError> {
let notification = serde_json::json!({
"jsonrpc": "2.0",
"method": method,
"params": params.unwrap_or(serde_json::Value::Null)
});
self.send_json(&notification).await
}
async fn send_json(&mut self, msg: &serde_json::Value) -> Result<(), AstrBotError> {
let content = serde_json::to_string(msg).map_err(|e| AstrBotError::Json(e))?;
let header = format!("Content-Length: {}\r\n\r\n", content.len());
if let Some(stdin) = &mut self.stdin {
stdin.write_all(header.as_bytes()).await?;
stdin.write_all(content.as_bytes()).await?;
stdin.flush().await?;
}
Ok(())
}
/// Read responses from LSP server
pub async fn read_responses(&mut self) -> Result<(), AstrBotError> {
if let Some(stdout) = &mut self.stdout {
let mut header_buf = Vec::new();
loop {
// Read header line
header_buf.clear();
stdout.read_until(b'\n', &mut header_buf).await?;
let header_str = String::from_utf8_lossy(&header_buf);
let content_length = header_str
.strip_prefix("Content-Length: ")
.and_then(|s| s.trim().parse::<usize>().ok())
.unwrap_or(0);
// Skip blank line (\r\n)
if header_buf.len() >= 2 && header_buf[0] == b'\r' && header_buf[1] == b'\n' {
// already at content
} else {
// Read the \r\n after Content-Length
let mut crlf = [0u8; 2];
stdout.read_exact(&mut crlf).await?;
}
// Read content
let mut content = vec![0u8; content_length];
stdout.read_exact(&mut content).await?;
let response: serde_json::Value =
serde_json::from_slice(&content).map_err(|e| AstrBotError::Json(e))?;
// Handle response
if let Some(id) = response.get("id").and_then(|v| v.as_i64()) {
if let Some(tx) = self.pending_requests.remove(&id) {
let _ = tx.send(response.to_string());
}
}
}
}
Ok(())
}
}
impl Default for LspClient {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl ProtocolClient for LspClient {
fn name(&self) -> &'static str {
"lsp"
}
fn is_connected(&self) -> bool {
self.connected
}
async fn connect(&mut self) -> Result<(), AstrBotError> {
self.connected = true;
tracing::debug!("ACP client connecting to unix socket:{}", socket_path);
tracing::debug!("LSP client connected");
Ok(())
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.connected = false;
// Try graceful shutdown first
if let Err(_) = self.send_notification("shutdown", None).await {
// Ignore errors during shutdown
}
if let Err(_) = self.send_notification("exit", None).await {
// Ignore errors during exit
}
// Terminate LSP server process
if let Some(mut child) = self.child.take() {
let _ = child.kill().await;
}
self.stdin = None;
self.stdout = None;
tracing::info!("LSP client disconnected");
Ok(())
}
}
// ============================================================================
// MCP Client - Model Context Protocol
// ============================================================================
pub struct McpClient {
connected: bool,
server_url: Option<String>,
transport: McpTransport,
}
enum McpTransport {
Stdio {
child: Option<Child>,
},
Http {
client: Option<reqwest::Client>,
base_url: String,
},
#[allow(dead_code)]
Sse {
client: Option<reqwest::Client>,
base_url: String,
},
}
impl McpClient {
pub fn new() -> Self {
Self {
connected: false,
server_url: None,
transport: McpTransport::Stdio { child: None },
}
}
pub fn set_connected(&mut self, connected: bool) {
self.connected = connected;
}
/// Connect to MCP server via stdio
pub async fn connect_to_stdio_server(
&mut self,
command: Vec<String>,
) -> Result<(), AstrBotError> {
if command.is_empty() {
return Err(AstrBotError::InvalidState("MCP command is empty".into()));
}
let mut cmd = Command::new(&command[0]);
if command.len() > 1 {
cmd.args(&command[1..]);
}
cmd.stdin(std::process::Stdio::piped());
cmd.stdout(std::process::Stdio::piped());
cmd.stderr(std::process::Stdio::piped());
let child = cmd.spawn().map_err(|e| {
AstrBotError::ConnectionFailed(format!("Failed to spawn MCP server: {e}"))
})?;
self.transport = McpTransport::Stdio { child: Some(child) };
self.connected = true;
self.server_url = Some("stdio".to_string());
tracing::info!("MCP client connected via stdio");
Ok(())
}
/// Connect to MCP server via HTTP/SSE
pub async fn connect_to_http_server(&mut self, base_url: &str) -> Result<(), AstrBotError> {
let client = reqwest::Client::new();
self.transport = McpTransport::Http {
client: Some(client),
base_url: base_url.to_string(),
};
self.connected = true;
self.server_url = Some(base_url.to_string());
tracing::info!("MCP client connected via HTTP: {}", base_url);
Ok(())
}
/// List available tools on MCP server
pub async fn list_tools(&self) -> Result<Vec<McpTool>, AstrBotError> {
match &self.transport {
McpTransport::Http { client, base_url } => {
if let Some(client) = client {
let url = format!("{}/tools", base_url);
let response = client
.get(&url)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(e.to_string()))?;
let tools: Vec<McpTool> = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(e.to_string()))?;
return Ok(tools);
}
}
McpTransport::Stdio { .. } => {
// Stdio-based tool listing would require JSON-RPC communication
}
McpTransport::Sse { .. } => {}
}
Ok(vec![])
}
/// Call a tool on MCP server
pub async fn call_tool(
&self,
tool_name: &str,
arguments: serde_json::Value,
) -> Result<serde_json::Value, AstrBotError> {
match &self.transport {
McpTransport::Http { client, base_url } => {
if let Some(client) = client {
let url = format!("{}/call", base_url);
let request = serde_json::json!({
"tool": tool_name,
"arguments": arguments
});
let response = client
.post(&url)
.json(&request)
.send()
.await
.map_err(|e| AstrBotError::ConnectionFailed(e.to_string()))?;
let result: serde_json::Value = response
.json()
.await
.map_err(|e| AstrBotError::Protocol(e.to_string()))?;
return Ok(result);
}
}
McpTransport::Stdio { child: _ } => {
// Stdio-based tool calling would require JSON-RPC communication
tracing::debug!("MCP stdio tool call: {} with {:?}", tool_name, arguments);
}
McpTransport::Sse { .. } => {}
}
Err(AstrBotError::NotConnected("MCP transport not ready".into()))
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpTool {
pub name: String,
pub description: String,
pub input_schema: serde_json::Value,
}
impl Default for McpClient {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl ProtocolClient for McpClient {
fn name(&self) -> &'static str {
"mcp"
}
fn is_connected(&self) -> bool {
self.connected
}
async fn connect(&mut self) -> Result<(), AstrBotError> {
self.connected = true;
tracing::debug!("MCP client connected");
Ok(())
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.connected = false;
if let McpTransport::Stdio { child } = &mut self.transport {
if let Some(mut child) = child.take() {
let _ = child.kill().await;
}
}
tracing::info!("MCP client disconnected");
Ok(())
}
}
// ============================================================================
// ACP Client - Agent Communication Protocol (e.g., Google A2A)
// ============================================================================
pub struct AcpClient {
connected: bool,
server_url: Option<String>,
tcp_stream: Option<TcpStream>,
pending_requests: HashMap<String, tokio::sync::oneshot::Sender<serde_json::Value>>,
request_id: u64,
}
impl AcpClient {
pub fn new() -> Self {
Self {
connected: false,
server_url: None,
tcp_stream: None,
pending_requests: HashMap::new(),
request_id: 0,
}
}
pub fn set_connected(&mut self, connected: bool) {
self.connected = connected;
}
/// Connect to ACP server via TCP
pub async fn connect_to_tcp(&mut self, host: &str, port: u16) -> Result<(), AstrBotError> {
let addr = format!("{}:{}", host, port);
let stream = TcpStream::connect(&addr)
.await
.map_err(|e| AstrBotError::ConnectionFailed(format!("TCP connect failed: {e}")))?;
self.tcp_stream = Some(stream);
self.connected = true;
self.server_url = Some(addr.clone());
tracing::info!("ACP client connected via TCP: {}", addr);
Ok(())
}
/// Connect to ACP server via Unix socket
/// Note: For now, Unix socket support is simplified - just mark as connected
pub async fn connect_to_unix_socket(&mut self, socket_path: &str) -> Result<(), AstrBotError> {
let _stream = UnixStream::connect(socket_path).await.map_err(|e| {
AstrBotError::ConnectionFailed(format!("Unix socket connect failed: {e}"))
})?;
self.connected = true;
self.server_url = Some(socket_path.to_string());
tracing::info!("ACP client connected via Unix socket: {}", socket_path);
Ok(())
}
/// Send a request and wait for response
pub async fn call_tool(
&mut self,
server_name: &str,
tool_name: &str,
arguments: serde_json::Value,
) -> Result<serde_json::Value, AstrBotError> {
let request_id = format!("{}-{}", self.request_id, server_name);
self.request_id += 1;
let (tx, rx) = tokio::sync::oneshot::channel();
self.pending_requests.insert(request_id.clone(), tx);
let request = serde_json::json!({
"jsonrpc": "2.0",
"id": request_id,
"method": format!("{}/{}", server_name, tool_name),
"params": arguments
});
self.send_json(&request).await?;
rx.await
.map_err(|_| AstrBotError::Timeout("ACP request timed out".into()))
}
/// Send a notification
pub async fn send_notification(
&mut self,
method: &str,
params: Option<serde_json::Value>,
) -> Result<(), AstrBotError> {
let notification = serde_json::json!({
"jsonrpc": "2.0",
"method": method,
"params": params.unwrap_or(serde_json::Value::Null)
});
self.send_json(&notification).await
}
async fn send_json(&mut self, msg: &serde_json::Value) -> Result<(), AstrBotError> {
let content = serde_json::to_string(msg).map_err(|e| AstrBotError::Json(e))?;
let header = serde_json::to_string(&serde_json::json!({
"content-length": content.len()
}))? + "\n";
if let Some(stream) = &mut self.tcp_stream {
stream.write_all(header.as_bytes()).await?;
stream.write_all(content.as_bytes()).await?;
}
Ok(())
}
/// Read messages from ACP server
pub async fn read_messages(&mut self) -> Result<(), AstrBotError> {
// Simplified - actual implementation would use BufReader
tracing::debug!("ACP read_messages called");
Ok(())
}
}
@@ -144,79 +587,25 @@ impl Default for AcpClient {
#[async_trait]
impl ProtocolClient for AcpClient {
async fn connect(&mut self) -> Result<(), AstrBotError> {
self.connected = true;
Ok(())
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.connected = false;
self.server_url = None;
Ok(())
fn name(&self) -> &'static str {
"acp"
}
fn is_connected(&self) -> bool {
self.connected
}
fn name(&self) -> &'static str {
"acp-client"
}
}
// ============================================================================
// ABP Client - AstrBot Protocol (internal stars)
// ============================================================================
#[derive(Debug, Default)]
pub struct AbpClient {
connected: bool,
stars: HashMap<String, String>,
}
impl AbpClient {
pub fn new() -> Self {
Self::default()
}
pub fn set_connected(&mut self, connected: bool) {
self.connected = connected;
}
pub fn register_star(&mut self, name: &str, _handler: &str) {
self.stars.insert(name.to_string(), name.to_string());
tracing::debug!("Star '{}' registered with ABP client", name);
}
pub fn unregister_star(&mut self, name: &str) {
self.stars.remove(name);
tracing::debug!("Star '{}' unregistered from ABP client", name);
}
pub fn list_stars(&self) -> Vec<String> {
self.stars.keys().cloned().collect()
}
}
#[async_trait]
impl ProtocolClient for AbpClient {
async fn connect(&mut self) -> Result<(), AstrBotError> {
self.connected = true;
tracing::debug!("ABP client connecting to internal stars");
tracing::debug!("ACP client connected");
Ok(())
}
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
self.connected = false;
self.stars.clear();
self.tcp_stream = None;
tracing::info!("ACP client disconnected");
Ok(())
}
fn is_connected(&self) -> bool {
self.connected
}
fn name(&self) -> &'static str {
"abp-client"
}
}

View File

@@ -1,18 +1,10 @@
//! Python bindings for AstrBot Core
use pyo3::prelude::*;
use pyo3::types::PyList;
/// Run the AstrBot Core CLI.
#[pyfunction]
pub fn cli(py: Python<'_>) -> PyResult<()> {
let sys = py.import("sys")?;
let argv: Bound<'_, PyAny> = sys.getattr("argv")?;
let argv_list = argv.downcast::<PyList>()?;
let args: Vec<String> = argv_list.iter()
.skip(1)
.filter_map(|item| item.extract::<String>().ok())
.collect();
pub fn cli(args: Vec<String>) -> PyResult<()> {
crate::cli::cli_with_args(&args)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}

606
rust/src/server.rs Normal file
View File

@@ -0,0 +1,606 @@
AstrBot/rust/src/server.rs
```rust
//! AstrBot HTTP/WebSocket Server
//!
//! High-performance async HTTP server with WebSocket support for real-time communication.
//! Implements JWT authentication and REST API endpoints.
use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, RwLock};
use tokio_tungstenite::{accept_async, tungstenite::Message};
use uuid::Uuid;
// ============================================================================
// Error Types
// ============================================================================
#[derive(Debug, thiserror::Error)]
pub enum ServerError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("WebSocket error: {0}")]
WebSocket(String),
#[error("JWT error: {0}")]
Jwt(String),
#[error("Not found: {0}")]
NotFound(String),
#[error("Unauthorized")]
Unauthorized,
#[error("Bad request: {0}")]
BadRequest(String),
}
impl From<tokio_tungstenite::tungstenite::Error> for ServerError {
fn from(err: tokio_tungstenite::tungstenite::Error) -> Self {
ServerError::WebSocket(err.to_string())
}
}
// ============================================================================
// JWT Authentication
// ============================================================================
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtClaims {
pub sub: String,
pub exp: u64,
pub iat: u64,
}
pub struct JwtAuth {
secret: String,
expiry_secs: u64,
}
impl JwtAuth {
pub fn new(secret: String, expiry_secs: u64) -> Self {
Self {
secret,
expiry_secs,
}
}
pub fn generate_token(&self, username: &str) -> Result<String, ServerError> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOODY)
.map_err(|e| ServerError::Jwt(e.to_string()))?
.as_secs();
let claims = JwtClaims {
sub: username.to_string(),
exp: now + self.expiry_secs,
iat: now,
};
let header = format!("{{\"alg\":\"HS256\",\"typ\":\"JWT\"}}");
let header_b64 = BASE64.encode(header.as_bytes());
let payload_b64 = BASE64.encode(serde_json::to_vec(&claims).map_err(|e| ServerError::Jwt(e.to_string()))?;
let signature = self.sign_jwt(&format!("{}.{}", header_b64, payload_b64));
Ok(format!("{}.{}.{}", header_b64, payload_b64, signature))
}
pub fn verify_token(&self, token: &str) -> Result<JwtClaims, ServerError> {
let parts: Vec<&str> = token.split('.').collect();
if parts.len() != 3 {
return Err(ServerError::Jwt("Invalid token format".to_string()));
}
let expected_sig = self.sign_jwt(&format!("{}.{}", parts[0], parts[1]));
if expected_sig != parts[2] {
return Err(ServerError::Jwt("Invalid signature".to_string()));
}
let payload = BASE64
.decode(parts[1])
.map_err(|e| ServerError::Jwt(e.to_string()))?;
let claims: JwtClaims =
serde_json::from_slice(&payload).map_err(|e| ServerError::Jwt(e.to_string()))?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_err(|e| ServerError::Jwt(e.to_string()))?
.as_secs();
if claims.exp < now {
return Err(ServerError::Jwt("Token expired".to_string()));
}
Ok(claims)
}
fn sign_jwt(&self, data: &str) -> String {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
data.hash(&mut hasher);
self.secret.hash(&mut hasher);
BASE64.encode(&hasher.finish().to_be_bytes())
}
}
// ============================================================================
// HTTP Response Types
// ============================================================================
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum ApiResponse {
#[serde(rename = "success")]
Success { data: serde_json::Value },
#[serde(rename = "error")]
Error { code: i32, message: String },
}
impl ApiResponse {
pub fn success<T: Serialize>(data: T) -> Self {
ApiResponse::Success {
data: serde_json::to_value(data).unwrap_or(serde_json::Value::Null),
}
}
pub fn error(code: i32, message: impl Into<String>) -> Self {
ApiResponse::Error {
code,
message: message.into(),
}
}
}
#[derive(Debug, Clone)]
pub struct HttpResponse {
pub status: u16,
pub status_text: &'static str,
pub headers: Vec<(String, String)>,
pub body: Vec<u8>,
}
impl HttpResponse {
pub fn new(status: u16, status_text: &'static str) -> Self {
Self {
status,
status_text,
headers: Vec::new(),
body: Vec::new(),
}
}
pub fn json(status: u16, status_text: &'static str, body: serde_json::Value) -> Self {
let body_bytes = serde_json::to_vec(&body).unwrap_or_default();
Self {
status,
status_text,
headers: vec![
("Content-Type".to_string(), "application/json".to_string()),
("Content-Length".to_string(), body_bytes.len().to_string()),
],
body: body_bytes,
}
}
pub fn html(status: u16, status_text: &'static str, body: impl Into<Vec<u8>>) -> Self {
let body = body.into();
Self {
status,
status_text,
headers: vec![
("Content-Type".to_string(), "text/html".to_string()),
("Content-Length".to_string(), body.len().to_string()),
],
body,
}
}
pub fn to_vec(&self) -> Vec<u8> {
let mut buf = format!("HTTP/1.1 {} {}\r\n", self.status, self.status_text).into_bytes();
for (key, value) in &self.headers {
buf.extend_from_slice(format!("{}: {}\r\n", key, value).as_bytes());
}
buf.extend_from_slice(b"\r\n");
buf.extend_from_slice(&self.body);
buf
}
}
// ============================================================================
// WebSocket Connection Manager
// ============================================================================
#[derive(Debug, Clone)]
pub struct WsConnection {
pub id: String,
pub authenticated: bool,
pub username: Option<String>,
}
#[derive(Default)]
pub struct WsManager {
connections: RwLock<HashMap<String, mpsc::Sender<String>>>,
}
impl WsManager {
pub fn new() -> Self {
Self {
connections: RwLock::new(HashMap::new()),
}
}
pub async fn broadcast(&self, message: &str) {
let connections = self.connections.read().await;
for sender in connections.values() {
let _ = sender.send(message.to_string()).await;
}
}
pub async fn send_to(&self, id: &str, message: &str) -> Result<(), ServerError> {
let connections = self.connections.read().await;
let sender = connections.get(id).ok_or_else(|| ServerError::NotFound(id.to_string()))?;
sender.send(message.to_string()).await.map_err(|_| {
ServerError::WebSocket("Failed to send".to_string())
})?;
Ok(())
}
pub async fn register(&self, id: String, sender: mpsc::Sender<String>) {
self.connections.write().await.insert(id, sender);
}
pub async fn unregister(&self, id: &str) {
self.connections.write().await.remove(id);
}
pub async fn count(&self) -> usize {
self.connections.read().await.len()
}
}
// ============================================================================
// HTTP Request
// ============================================================================
#[derive(Debug)]
pub struct HttpRequest {
pub method: String,
pub path: String,
pub headers: HashMap<String, String>,
pub body: Vec<u8>,
}
impl HttpRequest {
pub fn parse(buf: &[u8]) -> Result<Self, ServerError> {
let header_end = buf.windows(4).position(|w| w == b"\r\n\r\n");
let header_end = header_end.ok_or_else(|| ServerError::BadRequest("Invalid HTTP header".to_string()))?;
let header_str = String::from_utf8_lossy(&buf[..header_end]);
let body = buf[header_end + 4..].to_vec();
let mut lines = header_str.split("\r\n");
let request_line = lines.next().ok_or_else(|| ServerError::BadRequest("Missing request line".to_string()))?;
let parts: Vec<&str> = request_line.split_whitespace().collect();
if parts.len() < 2 {
return Err(ServerError::BadRequest("Invalid request line".to_string()));
}
let method = parts[0].to_string();
let path = parts[1].to_string();
let mut headers = HashMap::new();
for line in lines {
if let Some((key, value)) = line.split_once(':') {
headers.insert(key.trim().to_lowercase(), value.trim().to_string());
}
}
Ok(HttpRequest {
method,
path,
headers,
body,
})
}
pub fn get_header(&self, key: &str) -> Option<&str> {
self.headers.get(&key.to_lowercase()).map(|s| s.as_str())
}
pub fn auth_token(&self) -> Option<&str> {
self.get_header("authorization")
.and_then(|v| v.strip_prefix("Bearer "))
}
}
// ============================================================================
// Server State
// ============================================================================
#[derive(Clone)]
pub struct ServerState {
pub orchestrator: Arc<crate::Orchestrator>,
pub jwt_auth: Arc<JwtAuth>,
pub ws_manager: Arc<WsManager>,
pub dashboard_password: String,
}
impl ServerState {
pub fn new(jwt_secret: String, dashboard_password: String) -> Self {
Self {
orchestrator: Arc::new(crate::Orchestrator::new()),
jwt_auth: Arc::new(JwtAuth::new(jwt_secret, 86400)),
ws_manager: Arc::new(WsManager::new()),
dashboard_password,
}
}
}
// ============================================================================
// API Handlers
// ============================================================================
async fn handle_api(
state: &ServerState,
req: &HttpRequest,
) -> Result<HttpResponse, ServerError> {
let (path, method) = (req.path.as_str(), req.method.as_str());
// Login endpoint (no auth required)
if path == "/api/auth/login" && method == "POST" {
return handle_login(state, req).await;
}
// Authenticate other requests
let token = req.auth_token().ok_or(ServerError::Unauthorized)?;
state.jwt_auth.verify_token(token).map_err(|_| ServerError::Unauthorized)?;
match (path, method) {
("/api/stats", "GET") => {
let stats = state.orchestrator.stats();
Ok(HttpResponse::json(
200,
"OK",
serde_json::json!({
"message_count": stats.message_count(),
"uptime_seconds": stats.uptime_seconds(),
"last_activity": stats.last_activity_time(),
}),
))
}
("/api/sessions", "GET") => {
Ok(HttpResponse::json(
200,
"OK",
serde_json::json!({
"sessions": state.orchestrator.list_stars(),
}),
))
}
("/api/health", "GET") => {
Ok(HttpResponse::json(
200,
"OK",
serde_json::json!({
"status": "ok",
"orchestrator_running": state.orchestrator.is_running(),
"websocket_connections": state.ws_manager.count().await,
}),
))
}
_ if path.starts_with("/api/sessions/") && method == "DELETE" => {
let session_id = path.strip_prefix("/api/sessions/").unwrap_or("");
state.orchestrator.unregister_star(session_id).map_err(|_| {
ServerError::NotFound("Session not found".to_string())
})?;
Ok(HttpResponse::json(200, "OK", serde_json::json!({"success": true})))
}
_ if path.starts_with("/api/sessions/") && method == "GET" => {
let session_id = path.strip_prefix("/api/sessions/").unwrap_or("");
Ok(HttpResponse::json(
200,
"OK",
serde_json::json!({"session_id": session_id}),
))
}
_ => Ok(HttpResponse::json(
404,
"Not Found",
ApiResponse::error(404, "Endpoint not found"),
)),
}
}
async fn handle_login(state: &ServerState, req: &HttpRequest) -> Result<HttpResponse, ServerError> {
let body = String::from_utf8_lossy(&req.body);
let login: serde_json::Value =
serde_json::from_str(&body).map_err(|_| ServerError::BadRequest("Invalid JSON".to_string()))?;
let password = login
.get("password")
.and_then(|v| v.as_str())
.ok_or_else(|| ServerError::BadRequest("Missing password".to_string()))?;
if password != state.dashboard_password {
return Ok(HttpResponse::json(
401,
"Unauthorized",
ApiResponse::error(401, "Invalid credentials"),
));
}
let token = state.jwt_auth.generate_token("admin")?;
Ok(HttpResponse::json(
200,
"OK",
serde_json::json!({
"token": token,
"expires_in": 86400,
}),
))
}
// ============================================================================
// WebSocket Handler
// ============================================================================
async fn handle_websocket(state: ServerState, stream: TcpStream, addr: std::net::SocketAddr) {
let ws_stream = match accept_async(stream).await {
Ok(ws) => ws,
Err(e) => {
tracing::error!("WebSocket handshake failed from {}: {}", addr, e);
return;
}
};
let (writer, mut reader) = ws_stream.split();
let connection_id = Uuid::new_v4().to_string();
let (tx, rx) = mpsc::channel::<String>(100);
state.ws_manager.register(connection_id.clone(), tx).await;
// Send welcome message
let welcome = serde_json::json!({
"type": "connected",
"id": connection_id
});
if let Ok(text) = serde_json::to_string(&welcome) {
let _ = writer.send(Message::Text(text.into())).await;
}
// Handle messages
loop {
tokio::select! {
msg = reader.next() => {
match msg {
Some(Ok(Message::Text(text))) => {
if let Err(e) = handle_ws_message(&state, text.to_string()).await {
tracing::error!("WebSocket error: {}", e);
break;
}
}
Some(Ok(Message::Close(_))) | None => {
break;
}
_ => {}
}
}
Some(msg) = rx.recv() => {
if let Some(text) = msg {
if writer.send(Message::Text(text.into())).await.is_err() {
break;
}
}
}
}
}
state.ws_manager.unregister(&connection_id).await;
}
async fn handle_ws_message(state: &ServerState, msg: String) -> Result<(), ServerError> {
let parsed: serde_json::Value =
serde_json::from_str(&msg).map_err(|_| ServerError::BadRequest("Invalid JSON".to_string()))?;
let msg_type = parsed.get("type").and_then(|v| v.as_str()).unwrap_or("unknown");
let response = match msg_type {
"ping" => serde_json::json!({"type": "pong"}),
"subscribe" => serde_json::json!({"type": "subscribed"}),
"get_stats" => {
let stats = state.orchestrator.stats();
serde_json::json!({
"type": "stats",
"data": {
"message_count": stats.message_count(),
"uptime_seconds": stats.uptime_seconds(),
}
})
}
_ => serde_json::json!({"type": "unknown"}),
};
if !response.is_null() {
state
.ws_manager
.send_to(&connection_id, &serde_json::to_string(&response).unwrap())
.await
.ok();
}
Ok(())
}
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use futures_util::StreamExt;
// ============================================================================
// HTTP Server
// ============================================================================
pub struct HttpServer {
host: String,
port: u16,
state: ServerState,
}
impl HttpServer {
pub fn new(host: String, port: u16, state: ServerState) -> Self {
Self { host, port, state }
}
pub async fn serve(self) -> Result<(), ServerError> {
let addr = format!("{}:{}", self.host, self.port);
let listener = TcpListener::bind(&addr).await?;
tracing::info!("HTTP server listening on {}", addr);
loop {
match listener.accept().await {
Ok((stream, addr)) => {
let state = self.state.clone();
tokio::spawn(async move {
if let Err(e) = handle_connection(stream, state).await {
tracing::error!("Connection error from {}: {}", addr, e);
}
});
}
Err(e) => {
tracing::error!("Accept error: {}", e);
}
}
}
}
}
async fn handle_connection(stream: TcpStream, state: ServerState) -> Result<(), ServerError> {
let mut buf = vec![0u8; 8192];
let n = stream.read(&mut buf).await?;
buf.truncate(n);
let req = HttpRequest::parse(&buf)?;
// Check for WebSocket upgrade
if req.get_header("upgrade") == Some("websocket") {
// For now, return HTTP response indicating WebSocket not fully implemented
let response = HttpResponse::html(
426,
"Upgrade Required",
"<html><body>WebSocket upgrade not implemented</body></html>",
);
let mut stream = stream;
stream.write_all(&response.to_vec()).await?;
stream.flush().await?;
return Ok(());
}
let response = handle_api(&state, &req).await?;
let mut stream = stream;
stream.write_all(&response.to_vec()).await?;
stream.flush().await?;
Ok(())
}