Initial commit
This commit is contained in:
emmettlu
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
name = "agentos-tools"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
agentos-core.workspace = true
|
||||
agentos-kernel.workspace = true
|
||||
agentos-protocol.workspace = true
|
||||
async-trait.workspace = true
|
||||
hex.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,477 @@
|
||||
//! Fail-closed tool registry and approval contracts.
|
||||
|
||||
pub mod linux;
|
||||
|
||||
use agentos_core::{ExecutionIdentity, RunId};
|
||||
use agentos_protocol::{ConcurrencyClass, ToolCall, ToolSpec};
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ToolContext {
|
||||
pub identity: ExecutionIdentity,
|
||||
pub workspace: PathBuf,
|
||||
pub deadline: Instant,
|
||||
}
|
||||
|
||||
impl ToolContext {
|
||||
pub fn remaining(&self) -> Duration {
|
||||
self.deadline.saturating_duration_since(Instant::now())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait Tool: Send + Sync {
|
||||
fn spec(&self) -> ToolSpec;
|
||||
async fn execute(
|
||||
&self,
|
||||
context: &ToolContext,
|
||||
arguments: &Map<String, Value>,
|
||||
) -> Result<Value, ToolError>;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ToolResult {
|
||||
pub call_id: String,
|
||||
pub name: String,
|
||||
pub success: bool,
|
||||
pub output: String,
|
||||
pub duration_ms: u64,
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
pub struct ToolRegistry {
|
||||
handlers: BTreeMap<String, Arc<dyn Tool>>,
|
||||
allowed: BTreeSet<String>,
|
||||
output_limit_bytes: usize,
|
||||
argument_limit_bytes: usize,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for ToolRegistry {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter
|
||||
.debug_struct("ToolRegistry")
|
||||
.field("handlers", &self.handlers.keys().collect::<Vec<_>>())
|
||||
.field("allowed", &self.allowed)
|
||||
.field("output_limit_bytes", &self.output_limit_bytes)
|
||||
.field("argument_limit_bytes", &self.argument_limit_bytes)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolRegistry {
|
||||
pub fn new(
|
||||
allowed: impl IntoIterator<Item = String>,
|
||||
output_limit_bytes: usize,
|
||||
argument_limit_bytes: usize,
|
||||
) -> Result<Self, ToolError> {
|
||||
if output_limit_bytes == 0 || argument_limit_bytes == 0 {
|
||||
return Err(ToolError::InvalidLimits);
|
||||
}
|
||||
Ok(Self {
|
||||
handlers: BTreeMap::new(),
|
||||
allowed: allowed.into_iter().collect(),
|
||||
output_limit_bytes,
|
||||
argument_limit_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<(), ToolError> {
|
||||
let spec = tool.spec();
|
||||
spec.validate()?;
|
||||
if self.handlers.contains_key(&spec.name) {
|
||||
return Err(ToolError::DuplicateTool(spec.name));
|
||||
}
|
||||
self.handlers.insert(spec.name, tool);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn specs(&self) -> Vec<ToolSpec> {
|
||||
self.handlers
|
||||
.iter()
|
||||
.filter(|(name, _)| self.allowed.contains(*name))
|
||||
.map(|(_, tool)| tool.spec())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn concurrency_for(&self, name: &str) -> ConcurrencyClass {
|
||||
if !self.allowed.contains(name) {
|
||||
return ConcurrencyClass::Exclusive;
|
||||
}
|
||||
self.handlers
|
||||
.get(name)
|
||||
.map_or(ConcurrencyClass::Exclusive, |tool| tool.spec().concurrency)
|
||||
}
|
||||
|
||||
pub async fn execute(&self, context: &ToolContext, call: &ToolCall) -> ToolResult {
|
||||
let started = Instant::now();
|
||||
let result = self.execute_inner(context, call).await;
|
||||
let (success, output) = match result {
|
||||
Ok(Value::String(output)) => (true, output),
|
||||
Ok(output) => (
|
||||
true,
|
||||
serde_json::to_string(&output).unwrap_or_else(|error| {
|
||||
json!({"error": error.to_string(), "error_type": "serialization"}).to_string()
|
||||
}),
|
||||
),
|
||||
Err(error) => (
|
||||
false,
|
||||
json!({
|
||||
"error": error.to_string(),
|
||||
"error_type": error.kind(),
|
||||
})
|
||||
.to_string(),
|
||||
),
|
||||
};
|
||||
let (output, truncated) = truncate_utf8(output, self.output_limit_bytes);
|
||||
ToolResult {
|
||||
call_id: call.call_id.clone(),
|
||||
name: call.name.clone(),
|
||||
success,
|
||||
output,
|
||||
duration_ms: u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
|
||||
truncated,
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_inner(
|
||||
&self,
|
||||
context: &ToolContext,
|
||||
call: &ToolCall,
|
||||
) -> Result<Value, ToolError> {
|
||||
if context.remaining().is_zero() {
|
||||
return Err(ToolError::DeadlineExceeded);
|
||||
}
|
||||
if !self.allowed.contains(&call.name) {
|
||||
return Err(ToolError::NotAllowed(call.name.clone()));
|
||||
}
|
||||
let handler = self
|
||||
.handlers
|
||||
.get(&call.name)
|
||||
.ok_or_else(|| ToolError::NotAllowed(call.name.clone()))?;
|
||||
let arguments = serde_json::to_vec(&call.arguments)?;
|
||||
if arguments.len() > self.argument_limit_bytes {
|
||||
return Err(ToolError::ArgumentsTooLarge);
|
||||
}
|
||||
validate_json_schema(
|
||||
&Value::Object(call.arguments.clone()),
|
||||
&handler.spec().input_schema,
|
||||
"$",
|
||||
)?;
|
||||
handler.execute(context, &call.arguments).await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ApprovalGrant {
|
||||
pub approval_id: String,
|
||||
pub run_id: RunId,
|
||||
pub tool_name: String,
|
||||
pub arguments_sha256: String,
|
||||
pub approver_uid: u32,
|
||||
pub expires_at_unix_ms: u64,
|
||||
pub consumed: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ApprovalLedger {
|
||||
grants: Mutex<BTreeMap<String, ApprovalGrant>>,
|
||||
}
|
||||
|
||||
impl ApprovalLedger {
|
||||
pub fn issue(
|
||||
&self,
|
||||
run_id: &RunId,
|
||||
call: &ToolCall,
|
||||
approver_uid: u32,
|
||||
ttl: Duration,
|
||||
) -> Result<ApprovalGrant, ToolError> {
|
||||
let now = unix_millis()?;
|
||||
let arguments_sha256 = arguments_digest(&call.arguments)?;
|
||||
let approval_id = hex::encode(Sha256::digest(
|
||||
format!(
|
||||
"{}:{}:{}:{}:{}",
|
||||
run_id, call.call_id, call.name, arguments_sha256, now
|
||||
)
|
||||
.as_bytes(),
|
||||
));
|
||||
let grant = ApprovalGrant {
|
||||
approval_id: approval_id.clone(),
|
||||
run_id: run_id.clone(),
|
||||
tool_name: call.name.clone(),
|
||||
arguments_sha256,
|
||||
approver_uid,
|
||||
expires_at_unix_ms: now
|
||||
.saturating_add(u64::try_from(ttl.as_millis()).unwrap_or(u64::MAX)),
|
||||
consumed: false,
|
||||
};
|
||||
self.grants
|
||||
.lock()
|
||||
.map_err(|_| ToolError::ApprovalLedgerPoisoned)?
|
||||
.insert(approval_id, grant.clone());
|
||||
Ok(grant)
|
||||
}
|
||||
|
||||
pub fn consume(
|
||||
&self,
|
||||
approval_id: &str,
|
||||
run_id: &RunId,
|
||||
call: &ToolCall,
|
||||
caller_uid: u32,
|
||||
) -> Result<ApprovalGrant, ToolError> {
|
||||
let now = unix_millis()?;
|
||||
let digest = arguments_digest(&call.arguments)?;
|
||||
let mut grants = self
|
||||
.grants
|
||||
.lock()
|
||||
.map_err(|_| ToolError::ApprovalLedgerPoisoned)?;
|
||||
let grant = grants
|
||||
.get_mut(approval_id)
|
||||
.ok_or(ToolError::ApprovalMissing)?;
|
||||
if grant.consumed {
|
||||
return Err(ToolError::ApprovalReplayed);
|
||||
}
|
||||
if now > grant.expires_at_unix_ms {
|
||||
return Err(ToolError::ApprovalExpired);
|
||||
}
|
||||
if &grant.run_id != run_id
|
||||
|| grant.tool_name != call.name
|
||||
|| grant.arguments_sha256 != digest
|
||||
|| grant.approver_uid != caller_uid
|
||||
{
|
||||
return Err(ToolError::ApprovalMismatch);
|
||||
}
|
||||
grant.consumed = true;
|
||||
Ok(grant.clone())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn arguments_digest(arguments: &Map<String, Value>) -> Result<String, ToolError> {
|
||||
Ok(hex::encode(Sha256::digest(serde_json::to_vec(arguments)?)))
|
||||
}
|
||||
|
||||
pub fn validate_json_schema(value: &Value, schema: &Value, path: &str) -> Result<(), ToolError> {
|
||||
let schema = schema
|
||||
.as_object()
|
||||
.ok_or_else(|| ToolError::Schema(format!("invalid schema at {path}")))?;
|
||||
if let Some(expected) = schema.get("type").and_then(Value::as_str)
|
||||
&& !matches_type(value, expected)
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} must be {expected}")));
|
||||
}
|
||||
if let Some(allowed) = schema.get("enum").and_then(Value::as_array)
|
||||
&& !allowed.contains(value)
|
||||
{
|
||||
return Err(ToolError::Schema(format!(
|
||||
"{path} is not one of the allowed values"
|
||||
)));
|
||||
}
|
||||
if let Some(object) = value.as_object() {
|
||||
let properties = schema
|
||||
.get("properties")
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
if let Some(required) = schema.get("required").and_then(Value::as_array) {
|
||||
for field in required.iter().filter_map(Value::as_str) {
|
||||
if !object.contains_key(field) {
|
||||
return Err(ToolError::Schema(format!(
|
||||
"{path} is missing required field {field}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
if schema.get("additionalProperties") == Some(&Value::Bool(false)) {
|
||||
for field in object.keys() {
|
||||
if !properties.contains_key(field) {
|
||||
return Err(ToolError::Schema(format!(
|
||||
"{path} has unexpected field {field}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
for (field, child) in object {
|
||||
if let Some(child_schema) = properties.get(field) {
|
||||
validate_json_schema(child, child_schema, &format!("{path}.{field}"))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(array) = value.as_array() {
|
||||
if let Some(minimum) = schema.get("minItems").and_then(Value::as_u64)
|
||||
&& array.len() < usize::try_from(minimum).unwrap_or(usize::MAX)
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} has too few items")));
|
||||
}
|
||||
if let Some(maximum) = schema.get("maxItems").and_then(Value::as_u64)
|
||||
&& array.len() > usize::try_from(maximum).unwrap_or(usize::MAX)
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} has too many items")));
|
||||
}
|
||||
if let Some(child_schema) = schema.get("items") {
|
||||
for (index, child) in array.iter().enumerate() {
|
||||
validate_json_schema(child, child_schema, &format!("{path}[{index}]"))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(string) = value.as_str() {
|
||||
let length = string.chars().count();
|
||||
if let Some(minimum) = schema.get("minLength").and_then(Value::as_u64)
|
||||
&& length < usize::try_from(minimum).unwrap_or(usize::MAX)
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} is too short")));
|
||||
}
|
||||
if let Some(maximum) = schema.get("maxLength").and_then(Value::as_u64)
|
||||
&& length > usize::try_from(maximum).unwrap_or(usize::MAX)
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} is too long")));
|
||||
}
|
||||
}
|
||||
if let Some(number) = value.as_f64() {
|
||||
if let Some(minimum) = schema.get("minimum").and_then(Value::as_f64)
|
||||
&& number < minimum
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} is below the minimum")));
|
||||
}
|
||||
if let Some(maximum) = schema.get("maximum").and_then(Value::as_f64)
|
||||
&& number > maximum
|
||||
{
|
||||
return Err(ToolError::Schema(format!("{path} is above the maximum")));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn matches_type(value: &Value, expected: &str) -> bool {
|
||||
match expected {
|
||||
"null" => value.is_null(),
|
||||
"object" => value.is_object(),
|
||||
"array" => value.is_array(),
|
||||
"string" => value.is_string(),
|
||||
"boolean" => value.is_boolean(),
|
||||
"integer" => value.as_i64().is_some() || value.as_u64().is_some(),
|
||||
"number" => value.is_number(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_utf8(mut value: String, limit: usize) -> (String, bool) {
|
||||
const MARKER: &str = "\n{\"agentos_truncated\":true}";
|
||||
if value.len() <= limit {
|
||||
return (value, false);
|
||||
}
|
||||
let end = limit.saturating_sub(MARKER.len());
|
||||
let mut boundary = end.min(value.len());
|
||||
while !value.is_char_boundary(boundary) {
|
||||
boundary = boundary.saturating_sub(1);
|
||||
}
|
||||
value.truncate(boundary);
|
||||
value.push_str(&MARKER[..MARKER.len().min(limit.saturating_sub(value.len()))]);
|
||||
(value, true)
|
||||
}
|
||||
|
||||
fn unix_millis() -> Result<u64, ToolError> {
|
||||
let duration = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_err(|_| ToolError::ClockBeforeEpoch)?;
|
||||
Ok(u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ToolError {
|
||||
#[error("invalid tool registry byte limits")]
|
||||
InvalidLimits,
|
||||
#[error("duplicate tool: {0}")]
|
||||
DuplicateTool(String),
|
||||
#[error("tool is not registered or allowed: {0}")]
|
||||
NotAllowed(String),
|
||||
#[error("tool arguments exceed byte limit")]
|
||||
ArgumentsTooLarge,
|
||||
#[error("agent run deadline exceeded")]
|
||||
DeadlineExceeded,
|
||||
#[error("tool input failed validation: {0}")]
|
||||
Schema(String),
|
||||
#[error("tool execution failed: {0}")]
|
||||
Execution(String),
|
||||
#[error("tool serialization failed: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
#[error("tool protocol is invalid: {0}")]
|
||||
Protocol(#[from] agentos_protocol::ProtocolError),
|
||||
#[error("approval does not exist")]
|
||||
ApprovalMissing,
|
||||
#[error("approval has expired")]
|
||||
ApprovalExpired,
|
||||
#[error("approval has already been consumed")]
|
||||
ApprovalReplayed,
|
||||
#[error("approval does not match the run, tool, arguments, or approver")]
|
||||
ApprovalMismatch,
|
||||
#[error("approval ledger is unavailable")]
|
||||
ApprovalLedgerPoisoned,
|
||||
#[error("system clock is before the Unix epoch")]
|
||||
ClockBeforeEpoch,
|
||||
}
|
||||
|
||||
impl ToolError {
|
||||
const fn kind(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Schema(_) | Self::ArgumentsTooLarge => "tool_input",
|
||||
Self::NotAllowed(_) => "tool_policy",
|
||||
Self::DeadlineExceeded => "deadline",
|
||||
Self::ApprovalMissing
|
||||
| Self::ApprovalExpired
|
||||
| Self::ApprovalReplayed
|
||||
| Self::ApprovalMismatch => "approval",
|
||||
_ => "tool_execution",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use agentos_core::{AgentId, AgentPrincipal, SessionId};
|
||||
|
||||
#[test]
|
||||
fn approval_is_bound_to_exact_arguments_and_single_use() {
|
||||
let ledger = ApprovalLedger::default();
|
||||
let run_id = RunId::new();
|
||||
let call = ToolCall {
|
||||
call_id: "call-1".into(),
|
||||
name: "worldline.commit".into(),
|
||||
arguments: serde_json::from_value(json!({"id": "a"})).unwrap(),
|
||||
};
|
||||
let grant = ledger
|
||||
.issue(&run_id, &call, 1000, Duration::from_mins(1))
|
||||
.unwrap();
|
||||
ledger
|
||||
.consume(&grant.approval_id, &run_id, &call, 1000)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
ledger.consume(&grant.approval_id, &run_id, &call, 1000),
|
||||
Err(ToolError::ApprovalReplayed)
|
||||
));
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn context() -> ToolContext {
|
||||
ToolContext {
|
||||
identity: ExecutionIdentity {
|
||||
principal: AgentPrincipal {
|
||||
agent_id: AgentId::new(),
|
||||
owner_uid: 1000,
|
||||
agent_uid: 200_001,
|
||||
agent_gid: 200_001,
|
||||
},
|
||||
session_id: SessionId::new(),
|
||||
run_id: RunId::new(),
|
||||
},
|
||||
workspace: PathBuf::from("/tmp"),
|
||||
deadline: Instant::now() + Duration::from_secs(1),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
use crate::{Tool, ToolContext, ToolError, ToolRegistry};
|
||||
use agentos_kernel::{filesystem_usage, kernel_info, system_usage};
|
||||
use agentos_protocol::{ConcurrencyClass, ToolSpec};
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value, json};
|
||||
use std::collections::BTreeSet;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::process::Command;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct LinuxToolConfig {
|
||||
pub service_allowlist: BTreeSet<String>,
|
||||
pub output_limit_bytes: usize,
|
||||
pub argument_limit_bytes: usize,
|
||||
pub npu_device_root: PathBuf,
|
||||
pub npu_sysfs_root: PathBuf,
|
||||
pub npu_module_path: PathBuf,
|
||||
pub npu_proc_root: PathBuf,
|
||||
pub calrt_library: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl Default for LinuxToolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
service_allowlist: BTreeSet::new(),
|
||||
output_limit_bytes: 64 * 1024,
|
||||
argument_limit_bytes: 64 * 1024,
|
||||
npu_device_root: PathBuf::from("/dev"),
|
||||
npu_sysfs_root: PathBuf::from("/sys/class/calculet_chardev"),
|
||||
npu_module_path: PathBuf::from("/sys/module/calculet_pci"),
|
||||
npu_proc_root: PathBuf::from("/proc/calculet_monitor"),
|
||||
calrt_library: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn build_linux_registry(config: LinuxToolConfig) -> Result<ToolRegistry, ToolError> {
|
||||
let mut tools: Vec<Arc<dyn Tool>> = vec![
|
||||
Arc::new(SystemInspectTool),
|
||||
Arc::new(NpuInspectTool {
|
||||
device_root: config.npu_device_root,
|
||||
sysfs_root: config.npu_sysfs_root,
|
||||
module_path: config.npu_module_path,
|
||||
proc_root: config.npu_proc_root,
|
||||
calrt_library: config.calrt_library,
|
||||
}),
|
||||
Arc::new(PackageQueryTool),
|
||||
];
|
||||
if !config.service_allowlist.is_empty() {
|
||||
tools.push(Arc::new(ServiceInspectTool {
|
||||
allowlist: config.service_allowlist.clone(),
|
||||
}));
|
||||
tools.push(Arc::new(JournalReadTool {
|
||||
allowlist: config.service_allowlist,
|
||||
}));
|
||||
}
|
||||
let allowed = tools
|
||||
.iter()
|
||||
.map(|tool| tool.spec().name)
|
||||
.collect::<Vec<_>>();
|
||||
let mut registry = ToolRegistry::new(
|
||||
allowed,
|
||||
config.output_limit_bytes,
|
||||
config.argument_limit_bytes,
|
||||
)?;
|
||||
for tool in tools {
|
||||
registry.register(tool)?;
|
||||
}
|
||||
Ok(registry)
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct SystemInspectTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SystemInspectTool {
|
||||
fn spec(&self) -> ToolSpec {
|
||||
no_argument_spec(
|
||||
"system_inspect",
|
||||
"Read Linux kernel, memory, load, uptime, and root filesystem facts.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_context: &ToolContext,
|
||||
_arguments: &Map<String, Value>,
|
||||
) -> Result<Value, ToolError> {
|
||||
let kernel = kernel_info();
|
||||
let system = system_usage();
|
||||
let filesystem = filesystem_usage(Path::new("/"))
|
||||
.map_err(|error| ToolError::Execution(error.to_string()))?;
|
||||
let os_release = tokio::fs::read_to_string("/etc/os-release")
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
Ok(json!({
|
||||
"kernel": kernel,
|
||||
"system": system,
|
||||
"root_filesystem": filesystem,
|
||||
"os_release": parse_os_release(&os_release),
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct NpuInspectTool {
|
||||
pub device_root: PathBuf,
|
||||
pub sysfs_root: PathBuf,
|
||||
pub module_path: PathBuf,
|
||||
pub proc_root: PathBuf,
|
||||
pub calrt_library: Option<PathBuf>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
#[allow(clippy::struct_excessive_bools)]
|
||||
struct NpuStatus {
|
||||
devices: Vec<String>,
|
||||
driver_loaded: bool,
|
||||
sysfs_present: bool,
|
||||
monitor_present: bool,
|
||||
calrt_library: Option<String>,
|
||||
ready: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for NpuInspectTool {
|
||||
fn spec(&self) -> ToolSpec {
|
||||
no_argument_spec(
|
||||
"npu_inspect",
|
||||
"Read CALCULET NPU devices, PCI driver sysfs, monitor, and calrt availability without opening the device.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
_context: &ToolContext,
|
||||
_arguments: &Map<String, Value>,
|
||||
) -> Result<Value, ToolError> {
|
||||
let mut devices = Vec::new();
|
||||
if let Ok(mut entries) = tokio::fs::read_dir(&self.device_root).await {
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
let name = entry.file_name();
|
||||
let name = name.to_string_lossy();
|
||||
if let Some(suffix) = name.strip_prefix("calculet")
|
||||
&& !suffix.is_empty()
|
||||
&& suffix.bytes().all(|byte| byte.is_ascii_digit())
|
||||
{
|
||||
devices.push(entry.path().to_string_lossy().into_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
devices.sort();
|
||||
let calrt_library = self
|
||||
.calrt_library
|
||||
.as_ref()
|
||||
.filter(|path| path.is_file())
|
||||
.cloned()
|
||||
.or_else(find_calrt_library);
|
||||
let status = NpuStatus {
|
||||
driver_loaded: self.module_path.is_dir(),
|
||||
sysfs_present: self.sysfs_root.is_dir(),
|
||||
monitor_present: self.proc_root.exists(),
|
||||
ready: !devices.is_empty() && calrt_library.is_some(),
|
||||
devices,
|
||||
calrt_library: calrt_library.map(|path| path.to_string_lossy().into_owned()),
|
||||
};
|
||||
serde_json::to_value(status).map_err(ToolError::from)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ServiceInspectTool {
|
||||
pub allowlist: BTreeSet<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ServiceInspectTool {
|
||||
fn spec(&self) -> ToolSpec {
|
||||
ToolSpec {
|
||||
name: "service_inspect".into(),
|
||||
description: "Read one explicitly allowlisted systemd service state.".into(),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {"unit": {"type": "string", "maxLength": 128}},
|
||||
"required": ["unit"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
strict: true,
|
||||
concurrency: ConcurrencyClass::Safe,
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
context: &ToolContext,
|
||||
arguments: &Map<String, Value>,
|
||||
) -> Result<Value, ToolError> {
|
||||
let unit = string_argument(arguments, "unit")?;
|
||||
require_unit(unit, &self.allowlist)?;
|
||||
let output = run_read_only(
|
||||
"/usr/bin/systemctl",
|
||||
&[
|
||||
"show",
|
||||
unit,
|
||||
"--no-pager",
|
||||
"--property=Id,LoadState,ActiveState,SubState,UnitFileState,MainPID,ExecMainCode,ExecMainStatus,Result,StateChangeTimestamp",
|
||||
],
|
||||
context,
|
||||
)
|
||||
.await?;
|
||||
let properties = output
|
||||
.stdout
|
||||
.lines()
|
||||
.filter_map(|line| line.split_once('='))
|
||||
.map(|(key, value)| (key.to_owned(), Value::String(value.to_owned())))
|
||||
.collect::<Map<_, _>>();
|
||||
Ok(json!({
|
||||
"unit": unit,
|
||||
"exit_code": output.exit_code,
|
||||
"properties": properties,
|
||||
"stderr": output.stderr,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct JournalReadTool {
|
||||
pub allowlist: BTreeSet<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for JournalReadTool {
|
||||
fn spec(&self) -> ToolSpec {
|
||||
ToolSpec {
|
||||
name: "journal_read".into(),
|
||||
description:
|
||||
"Read a bounded number of journal lines for an allowlisted systemd service.".into(),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"unit": {"type": "string", "maxLength": 128},
|
||||
"lines": {"type": "integer", "minimum": 1, "maximum": 200},
|
||||
"since_minutes": {"type": "integer", "minimum": 1, "maximum": 1440}
|
||||
},
|
||||
"required": ["unit", "lines", "since_minutes"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
strict: true,
|
||||
concurrency: ConcurrencyClass::Safe,
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
context: &ToolContext,
|
||||
arguments: &Map<String, Value>,
|
||||
) -> Result<Value, ToolError> {
|
||||
let unit = string_argument(arguments, "unit")?;
|
||||
require_unit(unit, &self.allowlist)?;
|
||||
let lines = integer_argument(arguments, "lines")?.to_string();
|
||||
let since = format!("-{} minutes", integer_argument(arguments, "since_minutes")?);
|
||||
let output = run_read_only(
|
||||
"/usr/bin/journalctl",
|
||||
&[
|
||||
"--no-pager",
|
||||
"--output=short-iso",
|
||||
"--unit",
|
||||
unit,
|
||||
"--since",
|
||||
&since,
|
||||
"--lines",
|
||||
&lines,
|
||||
],
|
||||
context,
|
||||
)
|
||||
.await?;
|
||||
Ok(json!({
|
||||
"unit": unit,
|
||||
"exit_code": output.exit_code,
|
||||
"entries": output.stdout.lines().collect::<Vec<_>>(),
|
||||
"stderr": output.stderr,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PackageQueryTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for PackageQueryTool {
|
||||
fn spec(&self) -> ToolSpec {
|
||||
ToolSpec {
|
||||
name: "package_query".into(),
|
||||
description: "Query one exact Linux package without changing package state.".into(),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {"package": {"type": "string", "maxLength": 128}},
|
||||
"required": ["package"],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
strict: true,
|
||||
concurrency: ConcurrencyClass::Safe,
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
context: &ToolContext,
|
||||
arguments: &Map<String, Value>,
|
||||
) -> Result<Value, ToolError> {
|
||||
let package = string_argument(arguments, "package")?;
|
||||
if !valid_package(package) {
|
||||
return Err(ToolError::Schema("invalid package name".into()));
|
||||
}
|
||||
let (manager, binary, arguments): (&str, &str, Vec<&str>) =
|
||||
if Path::new("/usr/bin/dpkg-query").is_file() {
|
||||
(
|
||||
"dpkg",
|
||||
"/usr/bin/dpkg-query",
|
||||
vec![
|
||||
"--show",
|
||||
"--showformat=${binary:Package}\t${Version}\t${db:Status-Abbrev}\\n",
|
||||
package,
|
||||
],
|
||||
)
|
||||
} else if Path::new("/usr/bin/rpm").is_file() {
|
||||
(
|
||||
"rpm",
|
||||
"/usr/bin/rpm",
|
||||
vec![
|
||||
"-q",
|
||||
"--queryformat",
|
||||
"%{NAME}\t%{VERSION}-%{RELEASE}\t%{ARCH}\\n",
|
||||
package,
|
||||
],
|
||||
)
|
||||
} else if Path::new("/usr/bin/pacman").is_file() {
|
||||
("pacman", "/usr/bin/pacman", vec!["-Q", package])
|
||||
} else {
|
||||
return Err(ToolError::Execution(
|
||||
"no supported package query command is installed".into(),
|
||||
));
|
||||
};
|
||||
let output = run_read_only(binary, &arguments, context).await?;
|
||||
Ok(json!({
|
||||
"package": package,
|
||||
"manager": manager,
|
||||
"installed": output.exit_code == 0,
|
||||
"record": output.stdout,
|
||||
"stderr": output.stderr,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct CommandOutput {
|
||||
exit_code: i32,
|
||||
stdout: String,
|
||||
stderr: String,
|
||||
}
|
||||
|
||||
async fn run_read_only(
|
||||
program: &str,
|
||||
arguments: &[&str],
|
||||
context: &ToolContext,
|
||||
) -> Result<CommandOutput, ToolError> {
|
||||
let timeout = context.remaining().min(Duration::from_secs(30));
|
||||
if timeout.is_zero() {
|
||||
return Err(ToolError::DeadlineExceeded);
|
||||
}
|
||||
let mut command = Command::new(program);
|
||||
command
|
||||
.args(arguments)
|
||||
.env("LC_ALL", "C.UTF-8")
|
||||
.env_remove("AGENTOS_OPENAI_COMPATIBLE_API_KEY")
|
||||
.env_remove("OPENAI_API_KEY")
|
||||
.kill_on_drop(true);
|
||||
let output = tokio::time::timeout(timeout, command.output())
|
||||
.await
|
||||
.map_err(|_| ToolError::DeadlineExceeded)?
|
||||
.map_err(|error| ToolError::Execution(error.to_string()))?;
|
||||
Ok(CommandOutput {
|
||||
exit_code: output.status.code().unwrap_or(125),
|
||||
stdout: String::from_utf8_lossy(&output.stdout)
|
||||
.chars()
|
||||
.take(64 * 1024)
|
||||
.collect(),
|
||||
stderr: String::from_utf8_lossy(&output.stderr)
|
||||
.chars()
|
||||
.take(4096)
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
|
||||
fn no_argument_spec(name: &str, description: &str) -> ToolSpec {
|
||||
ToolSpec {
|
||||
name: name.into(),
|
||||
description: description.into(),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
"additionalProperties": false
|
||||
}),
|
||||
strict: true,
|
||||
concurrency: ConcurrencyClass::Safe,
|
||||
}
|
||||
}
|
||||
|
||||
fn string_argument<'a>(
|
||||
arguments: &'a Map<String, Value>,
|
||||
name: &str,
|
||||
) -> Result<&'a str, ToolError> {
|
||||
arguments
|
||||
.get(name)
|
||||
.and_then(Value::as_str)
|
||||
.ok_or_else(|| ToolError::Schema(format!("{name} must be a string")))
|
||||
}
|
||||
|
||||
fn integer_argument(arguments: &Map<String, Value>, name: &str) -> Result<u64, ToolError> {
|
||||
arguments
|
||||
.get(name)
|
||||
.and_then(Value::as_u64)
|
||||
.ok_or_else(|| ToolError::Schema(format!("{name} must be a positive integer")))
|
||||
}
|
||||
|
||||
fn require_unit(unit: &str, allowlist: &BTreeSet<String>) -> Result<(), ToolError> {
|
||||
let syntactically_valid = unit.ends_with(".service")
|
||||
&& unit.len() <= 128
|
||||
&& unit
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'.' | b'@' | b'-'));
|
||||
if !syntactically_valid || !allowlist.contains(unit) {
|
||||
return Err(ToolError::NotAllowed(format!("service unit {unit}")));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn valid_package(package: &str) -> bool {
|
||||
!package.is_empty()
|
||||
&& package.len() <= 128
|
||||
&& package.bytes().enumerate().all(|(index, byte)| {
|
||||
byte.is_ascii_alphanumeric()
|
||||
|| (index > 0 && matches!(byte, b'+' | b'.' | b'_' | b':' | b'@' | b'-'))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_os_release(content: &str) -> Map<String, Value> {
|
||||
let wanted = ["ID", "VERSION_ID", "PRETTY_NAME", "NAME"];
|
||||
content
|
||||
.lines()
|
||||
.filter_map(|line| line.split_once('='))
|
||||
.filter(|(key, _)| wanted.contains(key))
|
||||
.map(|(key, value)| {
|
||||
(
|
||||
key.to_owned(),
|
||||
Value::String(value.trim().trim_matches('"').chars().take(512).collect()),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn find_calrt_library() -> Option<PathBuf> {
|
||||
[
|
||||
"/usr/lib/libcalrt.so",
|
||||
"/usr/local/lib/libcalrt.so",
|
||||
"/opt/calculet/lib/libcalrt.so",
|
||||
]
|
||||
.into_iter()
|
||||
.map(PathBuf::from)
|
||||
.find(|path| path.is_file())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn package_and_service_names_are_fail_closed() {
|
||||
assert!(valid_package("linux-image-amd64"));
|
||||
assert!(!valid_package("--root=/tmp"));
|
||||
let allowlist = BTreeSet::from(["agentos-npud.service".to_owned()]);
|
||||
assert!(require_unit("agentos-npud.service", &allowlist).is_ok());
|
||||
assert!(require_unit("ssh.service", &allowlist).is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user