mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
feat(rust): initial Rust core runtime with CLI support
- Add Rust orchestrator with async bootstrap pattern - Implement CLI with clap (start, stats, health subcommands) - Add protocol stubs (LSP, MCP, ACP, ABP) - Python bindings via pyo3 (_core module) - Use maturin as build backend - Add tombi.toml for schema config
This commit is contained in:
8
.gitignore
vendored
8
.gitignore
vendored
@@ -70,3 +70,11 @@ GenieData/
|
||||
dashboard/warker.js
|
||||
dashboard/bun.lock
|
||||
.pua/
|
||||
|
||||
# Rust build artifacts
|
||||
rust/target/
|
||||
|
||||
# Build outputs
|
||||
dist/
|
||||
*.whl
|
||||
*.so
|
||||
|
||||
18
astrbot/_internal/runtime/rust/__init__.py
Normal file
18
astrbot/_internal/runtime/rust/__init__.py
Normal file
@@ -0,0 +1,18 @@
|
||||
import sys
|
||||
|
||||
try:
|
||||
from ._core import cli as _cli
|
||||
|
||||
def cli():
|
||||
if len(sys.argv) == 1:
|
||||
sys.argv.append("--help")
|
||||
return _cli()
|
||||
except ImportError:
|
||||
from click import echo
|
||||
|
||||
def cli():
|
||||
echo("""
|
||||
AstrBot CLI(rust) is not available.
|
||||
Developer: maturin dev
|
||||
User: uv run astrbot-rs
|
||||
""")
|
||||
16
astrbot/_internal/runtime/rust/_core.pyi
Normal file
16
astrbot/_internal/runtime/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() -> None: ...
|
||||
@@ -5,6 +5,7 @@ description = "Easy-to-use multi-platform LLM chatbot and development framework"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12,<3.14"
|
||||
|
||||
|
||||
keywords = ["Astrbot", "Astrbot Module", "Astrbot Plugin"]
|
||||
|
||||
dependencies = [
|
||||
@@ -72,6 +73,9 @@ dependencies = [
|
||||
"openai-agents>=0.12.5",
|
||||
"fastapi>=0.135.1",
|
||||
]
|
||||
[project.scripts]
|
||||
astrbot = "astrbot.cli.__main__:cli"
|
||||
astrbot-rs = "astrbot._internal.runtime._core:cli"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
@@ -84,8 +88,7 @@ dev = [
|
||||
"textual-dev>=1.8.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
astrbot = "astrbot.cli.__main__:cli"
|
||||
|
||||
|
||||
[tool.ruff]
|
||||
exclude = ["astrbot/core/utils/t2i/local_strategy.py", "astrbot/api/all.py", "tests"]
|
||||
@@ -124,14 +127,13 @@ exclude = ["dashboard", "node_modules", "dist", "data", "tests"]
|
||||
[tool.hatch.metadata]
|
||||
allow-direct-references = true
|
||||
|
||||
# Include bundled dashboard dist even though it is not tracked by VCS.
|
||||
[tool.hatch.build.targets.wheel]
|
||||
artifacts = ["astrbot/dashboard/dist/**"]
|
||||
|
||||
# Custom build hook: builds the Vue dashboard and copies dist into the package.
|
||||
[tool.hatch.build.hooks.custom]
|
||||
path = "scripts/hatch_build.py"
|
||||
# AstrBot Core Rust extension
|
||||
#:schema strict = false
|
||||
[tool.maturin]
|
||||
manifest-path = "rust/Cargo.toml"
|
||||
module-name = "astrbot._internal.runtime.rust._core"
|
||||
build-hook = "buildHooks.custom:pre_build_hook"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
requires = ["maturin>=1.12.0"]
|
||||
build-backend = "maturin"
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
[package]
|
||||
name = "astrbot-core"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
version = "4.25.0"
|
||||
edition = "2024"
|
||||
|
||||
[lib]
|
||||
name = "astrbot_core"
|
||||
name = "astrbot__core"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
@@ -14,14 +14,19 @@ python = ["pyo3"]
|
||||
[dependencies]
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
toml = "1.1"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
anyhow = "1"
|
||||
thiserror = "2"
|
||||
async-trait = "0.1"
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
|
||||
pyo3 = { version = "0.24", features = ["full"], optional = true }
|
||||
pyo3 = { version = "0.28", features = [
|
||||
"full",
|
||||
"extension-module",
|
||||
], optional = true }
|
||||
|
||||
[profile.release]
|
||||
opt-level = 3
|
||||
|
||||
1
rust/buildHooks/__init__.py
Normal file
1
rust/buildHooks/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
# Build hooks package
|
||||
162
rust/buildHooks/custom.py
Normal file
162
rust/buildHooks/custom.py
Normal file
@@ -0,0 +1,162 @@
|
||||
"""
|
||||
Custom maturin build hook for AstrBot.
|
||||
|
||||
This hook runs before the native extension is built.
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def get_root():
|
||||
"""Get project root directory."""
|
||||
return Path(__file__).parent.parent.parent
|
||||
|
||||
|
||||
def get_dashboard_dist():
|
||||
"""Get the dashboard dist source directory."""
|
||||
root = get_root()
|
||||
return root / "dashboard" / "dist"
|
||||
|
||||
|
||||
def get_dashboard_target():
|
||||
"""Get the target directory in the Python package."""
|
||||
root = get_root()
|
||||
return root / "astrbot" / "dashboard" / "dist"
|
||||
|
||||
|
||||
def build_dashboard():
|
||||
"""Build the Vue dashboard using npm."""
|
||||
root = get_root()
|
||||
dashboard_src = root / "dashboard"
|
||||
|
||||
if not dashboard_src.exists():
|
||||
print("[maturin-hook] dashboard/ directory not found, skipping dashboard build")
|
||||
return False
|
||||
|
||||
# Install node_modules if needed
|
||||
if not (dashboard_src / "node_modules").exists():
|
||||
print("[maturin-hook] Installing dashboard Node dependencies...")
|
||||
subprocess.run(
|
||||
["npm", "install"],
|
||||
cwd=dashboard_src,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
print("[maturin-hook] npm install output:", subprocess.run(
|
||||
["npm", "install"],
|
||||
cwd=dashboard_src,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
).stderr)
|
||||
|
||||
# Build dashboard
|
||||
print("[maturin-hook] Building Vue dashboard (npm run build)...")
|
||||
result = subprocess.run(
|
||||
["npm", "run", "build"],
|
||||
cwd=dashboard_src,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
print(f"[maturin-hook] npm build failed: {result.stderr}")
|
||||
return False
|
||||
print(f"[maturin-hook] npm build stdout: {result.stdout}")
|
||||
return True
|
||||
|
||||
|
||||
def copy_dashboard_dist():
|
||||
"""Copy dashboard dist to Python package."""
|
||||
root = get_root()
|
||||
dist_src = get_dashboard_dist()
|
||||
dist_target = get_dashboard_target()
|
||||
|
||||
if not dist_src.exists():
|
||||
print("[maturin-hook] No dashboard/dist found after build")
|
||||
return False
|
||||
|
||||
# Remove existing target
|
||||
if dist_target.exists():
|
||||
shutil.rmtree(dist_target)
|
||||
|
||||
# Copy dist
|
||||
shutil.copytree(dist_src, dist_target)
|
||||
print(f"[maturin-hook] Dashboard dist copied to {dist_target}")
|
||||
return True
|
||||
|
||||
|
||||
def generate_placeholder():
|
||||
"""Generate placeholder if no dashboard."""
|
||||
root = get_root()
|
||||
dist_target = get_dashboard_target()
|
||||
dist_target.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
placeholder = dist_target / ".placeholder"
|
||||
if not any(dist_target.iterdir()):
|
||||
placeholder.write_text("# dashboard placeholder\n")
|
||||
print(f"[maturin-hook] Created placeholder at {placeholder}")
|
||||
|
||||
|
||||
class MaturinBuildHook:
|
||||
"""Maturin build hook class."""
|
||||
|
||||
def __init__(self, args):
|
||||
self.args = args
|
||||
print(f"[maturin-hook] MaturinBuildHook initialized with args: {args}")
|
||||
|
||||
def pre_build(self):
|
||||
"""Called before build."""
|
||||
print("[maturin-hook] pre_build called!")
|
||||
print(f"[maturin-hook] Current dir: {os.getcwd()}")
|
||||
print(f"[maturin-hook] sys.argv: {sys.argv}")
|
||||
print(f"[maturin-hook] ASTRBOT_BUILD_DASHBOARD: {os.environ.get('ASTRBOT_BUILD_DASHBOARD', 'NOT SET')}")
|
||||
|
||||
root = get_root()
|
||||
print(f"[maturin-hook] Project root: {root}")
|
||||
|
||||
if os.environ.get("ASTRBOT_BUILD_DASHBOARD", "").strip() == "1":
|
||||
print("[maturin-hook] Building dashboard...")
|
||||
if build_dashboard():
|
||||
copy_dashboard_dist()
|
||||
else:
|
||||
print("[maturin-hook] Skipping dashboard build (ASTRBOT_BUILD_DASHBOARD != 1)")
|
||||
generate_placeholder()
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def pre_build_hook(ctx):
|
||||
"""Entry point for maturin pre-build hook.
|
||||
|
||||
Args:
|
||||
ctx: BuildContext with manifest_path and cargo_args
|
||||
|
||||
Returns:
|
||||
int: 0 on success
|
||||
"""
|
||||
print("[maturin-hook] pre_build_hook called!")
|
||||
print(f"[maturin-hook] ctx: {ctx}")
|
||||
print(f"[maturin-hook] sys.argv: {sys.argv}")
|
||||
|
||||
root = Path(__file__).parent.parent.parent
|
||||
print(f"[maturin-hook] Project root: {root}")
|
||||
|
||||
if os.environ.get("ASTRBOT_BUILD_DASHBOARD", "").strip() == "1":
|
||||
print("[maturin-hook] Building dashboard...")
|
||||
if build_dashboard():
|
||||
copy_dashboard_dist()
|
||||
else:
|
||||
print("[maturin-hook] Skipping dashboard build (ASTRBOT_BUILD_DASHBOARD != 1)")
|
||||
generate_placeholder()
|
||||
|
||||
# Ensure target directory exists for wheel artifacts
|
||||
dist_target = get_dashboard_target()
|
||||
dist_target.mkdir(parents=True, exist_ok=True)
|
||||
if not any(dist_target.iterdir()):
|
||||
generate_placeholder()
|
||||
|
||||
return 0
|
||||
70
rust/src/cli.rs
Normal file
70
rust/src/cli.rs
Normal file
@@ -0,0 +1,70 @@
|
||||
//! AstrBot Core CLI
|
||||
|
||||
use anyhow::Result;
|
||||
use clap::{Parser, Subcommand};
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "astrbot-rs")]
|
||||
#[command(version = "0.1.0")]
|
||||
pub struct Cli {
|
||||
#[command(subcommand)]
|
||||
pub command: Commands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
pub enum Commands {
|
||||
/// Start the AstrBot core runtime
|
||||
Start {
|
||||
/// Host to bind to
|
||||
#[arg(long, default_value = "127.0.0.1")]
|
||||
host: String,
|
||||
/// Port to listen on
|
||||
#[arg(long, default_value_t = 8765)]
|
||||
port: u16,
|
||||
},
|
||||
/// Show runtime statistics
|
||||
Stats,
|
||||
/// Run health check
|
||||
Health,
|
||||
}
|
||||
|
||||
pub fn cli() -> Result<()> {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
cli_with_args(&args[1..])
|
||||
}
|
||||
|
||||
pub fn cli_with_args(args: &[String]) -> Result<()> {
|
||||
let cli = Cli::try_parse_from(args).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
|
||||
match cli.command {
|
||||
Commands::Start { host, port } => {
|
||||
start_runtime(&host, port)?;
|
||||
}
|
||||
Commands::Stats => {
|
||||
show_stats()?;
|
||||
}
|
||||
Commands::Health => {
|
||||
health_check()?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn start_runtime(host: &str, port: u16) -> Result<()> {
|
||||
println!("Starting AstrBot Core runtime on {host}:{port}");
|
||||
println!("AstrBot Core v{} is running", env!("CARGO_PKG_VERSION"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn show_stats() -> Result<()> {
|
||||
println!("AstrBot Core Statistics");
|
||||
println!("========================");
|
||||
println!("Version: {}", env!("CARGO_PKG_VERSION"));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn health_check() -> Result<()> {
|
||||
println!("OK");
|
||||
Ok(())
|
||||
}
|
||||
132
rust/src/config.rs
Normal file
132
rust/src/config.rs
Normal file
@@ -0,0 +1,132 @@
|
||||
//! Configuration management for AstrBot Core
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// AstrBot Core configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Config {
|
||||
/// Runtime settings
|
||||
pub runtime: RuntimeConfig,
|
||||
/// Protocol client settings
|
||||
pub protocols: ProtocolsConfig,
|
||||
/// Logging settings
|
||||
pub logging: LoggingConfig,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RuntimeConfig {
|
||||
/// Whether the runtime is in debug mode
|
||||
pub debug: bool,
|
||||
/// Data directory path
|
||||
pub data_dir: PathBuf,
|
||||
/// Config directory path
|
||||
pub config_dir: PathBuf,
|
||||
/// Temporary directory path
|
||||
pub temp_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl Default for RuntimeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
debug: false,
|
||||
data_dir: PathBuf::from("~/.astrbot/data"),
|
||||
config_dir: PathBuf::from("~/.astrbot/config"),
|
||||
temp_dir: PathBuf::from("~/.astrbot/temp"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ProtocolsConfig {
|
||||
/// LSP protocol settings
|
||||
pub lsp: ProtocolSettings,
|
||||
/// MCP protocol settings
|
||||
pub mcp: ProtocolSettings,
|
||||
/// ACP protocol settings
|
||||
pub acp: ProtocolSettings,
|
||||
/// ABP protocol settings
|
||||
pub abp: ProtocolSettings,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ProtocolSettings {
|
||||
/// Whether this protocol is enabled
|
||||
pub enabled: bool,
|
||||
/// Connection timeout in seconds
|
||||
pub timeout_secs: u64,
|
||||
/// Retry settings
|
||||
pub retry: RetrySettings,
|
||||
}
|
||||
|
||||
impl Default for ProtocolSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: true,
|
||||
timeout_secs: 30,
|
||||
retry: RetrySettings::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RetrySettings {
|
||||
/// Maximum number of retries
|
||||
pub max_retries: u32,
|
||||
/// Initial backoff delay in milliseconds
|
||||
pub initial_delay_ms: u64,
|
||||
/// Maximum backoff delay in milliseconds
|
||||
pub max_delay_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for RetrySettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_retries: 3,
|
||||
initial_delay_ms: 100,
|
||||
max_delay_ms: 5000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LoggingConfig {
|
||||
/// Log level (trace, debug, info, warn, error)
|
||||
pub level: String,
|
||||
/// Whether to use structured logging
|
||||
pub structured: bool,
|
||||
/// Log file path
|
||||
pub file: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl Default for LoggingConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
level: "info".to_string(),
|
||||
structured: true,
|
||||
file: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Load configuration from a file
|
||||
pub fn load(path: &PathBuf) -> anyhow::Result<Self> {
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let config: Config = toml::from_str(&content)?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Save configuration to a file
|
||||
pub fn save(&self, path: &PathBuf) -> anyhow::Result<()> {
|
||||
let content = toml::to_string_pretty(self)?;
|
||||
std::fs::write(path, content)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -18,11 +18,21 @@
|
||||
#![allow(clippy::cast_possible_truncation)]
|
||||
#![allow(clippy::cast_lossless)]
|
||||
#![allow(clippy::unnecessary_struct_initialization)]
|
||||
#![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 cli;
|
||||
pub mod error;
|
||||
pub mod orchestrator;
|
||||
pub mod message;
|
||||
pub mod stats;
|
||||
pub mod protocol;
|
||||
pub mod config;
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
pub mod python;
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
fn main() {
|
||||
println!("Hello, world!");
|
||||
}
|
||||
@@ -13,9 +13,10 @@ pub struct Message {
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum MessageType {
|
||||
#[default]
|
||||
Text,
|
||||
Image,
|
||||
Audio,
|
||||
@@ -25,12 +26,6 @@ pub enum MessageType {
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl Default for MessageType {
|
||||
fn default() -> Self {
|
||||
Self::Text
|
||||
}
|
||||
}
|
||||
|
||||
impl Message {
|
||||
#[must_use]
|
||||
pub fn new(id: String, content: String, sender: String) -> Self {
|
||||
|
||||
@@ -3,7 +3,9 @@
|
||||
use crate::error::AstrBotError;
|
||||
use crate::stats::RuntimeStats;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::RwLock;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use tokio::sync::broadcast;
|
||||
use tokio::time::{interval, Duration};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProtocolStatus {
|
||||
@@ -20,36 +22,121 @@ impl Default for ProtocolStatus {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
/// Main orchestrator coordinating all protocol clients
|
||||
pub struct Orchestrator {
|
||||
running: RwLock<bool>,
|
||||
shutdown_tx: Arc<RwLock<Option<broadcast::Sender<()>>>>,
|
||||
stars: RwLock<HashMap<String, String>>,
|
||||
stats: RuntimeStats,
|
||||
protocol_lsp: RwLock<ProtocolStatus>,
|
||||
protocol_mcp: RwLock<ProtocolStatus>,
|
||||
protocol_acp: RwLock<ProtocolStatus>,
|
||||
protocol_abp: RwLock<ProtocolStatus>,
|
||||
stats: RwLock<RuntimeStats>,
|
||||
}
|
||||
|
||||
impl Default for Orchestrator {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Orchestrator {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
Self {
|
||||
running: RwLock::new(false),
|
||||
shutdown_tx: Arc::new(RwLock::new(None)),
|
||||
stars: RwLock::new(HashMap::new()),
|
||||
stats: RwLock::new(RuntimeStats::default()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Start the orchestrator
|
||||
pub fn start(&self) -> Result<(), AstrBotError> {
|
||||
let mut running = self.running.write().map_err(|_| {
|
||||
AstrBotError::InvalidState("Failed to acquire write lock".into())
|
||||
})?;
|
||||
*running = true;
|
||||
{
|
||||
let mut running = self.running.write().map_err(|_| {
|
||||
AstrBotError::InvalidState("Failed to acquire write lock".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())
|
||||
})?;
|
||||
*shutdown_tx = Some(tx);
|
||||
}
|
||||
|
||||
tracing::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;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Main event loop
|
||||
pub async fn run_loop(&self) -> Result<(), AstrBotError> {
|
||||
if !self.is_running() {
|
||||
return Err(AstrBotError::InvalidState("Orchestrator not started".into()));
|
||||
}
|
||||
|
||||
tracing::info!("Orchestrator event loop started");
|
||||
let mut tick_interval = interval(Duration::from_secs(5));
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = tick_interval.tick() => {
|
||||
self.periodic_health_check();
|
||||
}
|
||||
_ = self.wait_for_shutdown() => {
|
||||
tracing::info!("Orchestrator shutdown signal received");
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if !self.is_running() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Orchestrator event loop stopped");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
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());
|
||||
|
||||
if let Some(tx) = tx {
|
||||
let mut rx = tx.subscribe();
|
||||
let _ = rx.recv().await;
|
||||
} else {
|
||||
tokio::time::sleep(Duration::MAX).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn periodic_health_check(&self) {
|
||||
tracing::debug!("Orchestrator health check");
|
||||
}
|
||||
|
||||
/// Stop the orchestrator
|
||||
pub fn stop(&self) -> Result<(), AstrBotError> {
|
||||
let mut running = self.running.write().map_err(|_| {
|
||||
AstrBotError::InvalidState("Failed to acquire write lock".into())
|
||||
})?;
|
||||
*running = false;
|
||||
{
|
||||
let mut running = self.running.write().map_err(|_| {
|
||||
AstrBotError::InvalidState("Failed to acquire write lock".into())
|
||||
})?;
|
||||
*running = false;
|
||||
}
|
||||
|
||||
if let Ok(tx_guard) = self.shutdown_tx.read() {
|
||||
if let Some(tx) = tx_guard.as_ref() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!("Orchestrator stopped");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -83,38 +170,25 @@ impl Orchestrator {
|
||||
}
|
||||
|
||||
pub fn record_activity(&self) {
|
||||
self.stats.record_message();
|
||||
if let Ok(stats) = self.stats.write() {
|
||||
stats.record_message();
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn stats(&self) -> RuntimeStats {
|
||||
self.stats.clone()
|
||||
self.stats.read().map(|s| s.clone()).unwrap_or_default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn get_protocol_status(&self, protocol: &str) -> Option<ProtocolStatus> {
|
||||
match protocol {
|
||||
"lsp" => self.protocol_lsp.read().ok().map(|p| p.clone()),
|
||||
"mcp" => self.protocol_mcp.read().ok().map(|p| p.clone()),
|
||||
"acp" => self.protocol_acp.read().ok().map(|p| p.clone()),
|
||||
"abp" => self.protocol_abp.read().ok().map(|p| p.clone()),
|
||||
_ => None,
|
||||
}
|
||||
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> {
|
||||
let status = match protocol {
|
||||
"lsp" => &self.protocol_lsp,
|
||||
"mcp" => &self.protocol_mcp,
|
||||
"acp" => &self.protocol_acp,
|
||||
"abp" => &self.protocol_abp,
|
||||
_ => return Err(AstrBotError::InvalidState(format!("Unknown protocol: {protocol}"))),
|
||||
};
|
||||
|
||||
let mut lock = status.write().map_err(|_| {
|
||||
AstrBotError::InvalidState("Failed to acquire write lock".into())
|
||||
})?;
|
||||
lock.connected = connected;
|
||||
pub fn set_protocol_connected(&self, _protocol: &str, _connected: bool) -> Result<(), AstrBotError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
222
rust/src/protocol.rs
Normal file
222
rust/src/protocol.rs
Normal file
@@ -0,0 +1,222 @@
|
||||
//! Protocol client implementations for ACP and ABP
|
||||
|
||||
use crate::error::AstrBotError;
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// 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;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// LSP Client
|
||||
// ============================================================================
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct LspClient {
|
||||
connected: bool,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_connected(&mut self, connected: bool) {
|
||||
self.connected = connected;
|
||||
}
|
||||
|
||||
pub async fn connect_to_server(&mut self, host: &str, port: u16) -> Result<(), AstrBotError> {
|
||||
self.server_url = Some(format!("{}:{}", host, port));
|
||||
self.connected = true;
|
||||
tracing::debug!("ACP client connecting to {}:{}", host, port);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn connect_to_unix_socket(&mut self, socket_path: &str) -> Result<(), AstrBotError> {
|
||||
self.server_url = Some(format!("unix://{}", socket_path));
|
||||
self.connected = true;
|
||||
tracing::debug!("ACP client connecting to unix socket:{}", socket_path);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AcpClient {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[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 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");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn disconnect(&mut self) -> Result<(), AstrBotError> {
|
||||
self.connected = false;
|
||||
self.stars.clear();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_connected(&self) -> bool {
|
||||
self.connected
|
||||
}
|
||||
|
||||
fn name(&self) -> &'static str {
|
||||
"abp-client"
|
||||
}
|
||||
}
|
||||
@@ -1,96 +1,25 @@
|
||||
//! Python bindings for AstrBot Core
|
||||
|
||||
use crate::orchestrator::Orchestrator;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::sync::GILOnceCell;
|
||||
use pyo3::types::PyDict;
|
||||
|
||||
static ORCHESTRATOR: GILOnceCell<Py<PythonOrchestrator>> = GILOnceCell::new();
|
||||
|
||||
#[pyclass]
|
||||
pub struct PythonOrchestrator {
|
||||
inner: Orchestrator,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PythonOrchestrator {
|
||||
#[new]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Orchestrator::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn start(&self) -> PyResult<()> {
|
||||
self.inner.start().map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
|
||||
}
|
||||
|
||||
pub fn stop(&self) -> PyResult<()> {
|
||||
self.inner.stop().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()
|
||||
}
|
||||
|
||||
#[allow(unsafe_code)]
|
||||
pub fn get_stats(&self) -> PyResult<Py<PyAny>> {
|
||||
let stats = self.inner.stats();
|
||||
let py = unsafe { Python::assume_gil_acquired() };
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("message_count", stats.message_count())?;
|
||||
dict.set_item("uptime_seconds", stats.uptime_seconds())?;
|
||||
if let Some(last) = stats.last_activity_time() {
|
||||
dict.set_item("last_activity", last)?;
|
||||
}
|
||||
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()))
|
||||
}
|
||||
|
||||
#[allow(unsafe_code)]
|
||||
pub fn get_protocol_status(&self, protocol: &str) -> Option<Py<PyAny>> {
|
||||
let status = self.inner.get_protocol_status(protocol)?;
|
||||
let py = unsafe { Python::assume_gil_acquired() };
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("connected", status.connected).ok()?;
|
||||
dict.set_item("name", status.name).ok()?;
|
||||
Some(dict.into())
|
||||
}
|
||||
}
|
||||
use pyo3::types::PyList;
|
||||
|
||||
/// Run the AstrBot Core CLI.
|
||||
#[pyfunction]
|
||||
pub fn get_orchestrator(py: Python<'_>) -> PyResult<&'static Py<PythonOrchestrator>> {
|
||||
if ORCHESTRATOR.get(py).is_none() {
|
||||
let orchestrator = Py::new(py, PythonOrchestrator::new())
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("Failed to create: {e}")))?;
|
||||
ORCHESTRATOR.set(py, orchestrator)
|
||||
.map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Already initialized"))?;
|
||||
}
|
||||
Ok(ORCHESTRATOR.get(py).expect("orchestrator should be initialized"))
|
||||
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();
|
||||
crate::cli::cli_with_args(&args)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
|
||||
}
|
||||
|
||||
/// Python module for AstrBot Core.
|
||||
#[pymodule]
|
||||
pub fn astrbot_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PythonOrchestrator>()?;
|
||||
m.add_function(wrap_pyfunction!(get_orchestrator, m)?)?;
|
||||
pub fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(cli, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
2
tombi.toml
Normal file
2
tombi.toml
Normal file
@@ -0,0 +1,2 @@
|
||||
[schema]
|
||||
strict = false
|
||||
Reference in New Issue
Block a user