mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
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:
16
astrbot/rust/_core.pyi
Normal file
16
astrbot/rust/_core.pyi
Normal 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
1850
rust/Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -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
464
rust/src/a2a.rs
Normal 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
726
rust/src/abp.rs
Normal 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
|
||||
}
|
||||
}
|
||||
@@ -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 } => {
|
||||
|
||||
1174
rust/src/config.rs
1174
rust/src/config.rs
File diff suppressed because it is too large
Load Diff
@@ -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),
|
||||
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(¬ification).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(¬ification).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"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
606
rust/src/server.rs
Normal 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(())
|
||||
}
|
||||
Reference in New Issue
Block a user