hehe_tools/builtin/
system.rs

1use crate::error::Result;
2use crate::traits::{Tool, ToolOutput};
3use async_trait::async_trait;
4use hehe_core::{Context, ToolDefinition};
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::env;
8
9pub struct GetSystemInfoTool {
10    def: ToolDefinition,
11}
12
13impl GetSystemInfoTool {
14    pub fn new() -> Self {
15        let def = ToolDefinition::new("get_system_info", "Get information about the current system");
16        Self { def }
17    }
18}
19
20impl Default for GetSystemInfoTool {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26#[derive(Serialize, Deserialize)]
27struct SystemInfo {
28    os: OsInfo,
29    process: ProcessInfo,
30    env: EnvInfo,
31}
32
33#[derive(Serialize, Deserialize)]
34struct OsInfo {
35    name: String,
36    arch: String,
37    family: String,
38}
39
40#[derive(Serialize, Deserialize)]
41struct ProcessInfo {
42    current_dir: Option<String>,
43    exe_path: Option<String>,
44    pid: u32,
45}
46
47#[derive(Serialize, Deserialize)]
48struct EnvInfo {
49    home: Option<String>,
50    user: Option<String>,
51    path: Option<String>,
52}
53
54#[async_trait]
55impl Tool for GetSystemInfoTool {
56    fn definition(&self) -> &ToolDefinition {
57        &self.def
58    }
59
60    async fn execute(&self, _ctx: &Context, _input: Value) -> Result<ToolOutput> {
61        let info = SystemInfo {
62            os: OsInfo {
63                name: env::consts::OS.to_string(),
64                arch: env::consts::ARCH.to_string(),
65                family: env::consts::FAMILY.to_string(),
66            },
67            process: ProcessInfo {
68                current_dir: env::current_dir()
69                    .ok()
70                    .map(|p| p.to_string_lossy().to_string()),
71                exe_path: env::current_exe()
72                    .ok()
73                    .map(|p| p.to_string_lossy().to_string()),
74                pid: std::process::id(),
75            },
76            env: EnvInfo {
77                home: env::var("HOME").ok().or_else(|| env::var("USERPROFILE").ok()),
78                user: env::var("USER").ok().or_else(|| env::var("USERNAME").ok()),
79                path: env::var("PATH").ok(),
80            },
81        };
82
83        ToolOutput::json(&info)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[tokio::test]
92    async fn test_get_system_info() {
93        let tool = GetSystemInfoTool::new();
94        let ctx = Context::new();
95
96        let output = tool.execute(&ctx, Value::Null).await.unwrap();
97        assert!(!output.is_error);
98
99        let info: SystemInfo = serde_json::from_str(&output.content).unwrap();
100        assert!(!info.os.name.is_empty());
101        assert!(!info.os.arch.is_empty());
102        assert!(info.process.pid > 0);
103    }
104
105    #[test]
106    fn test_definition() {
107        let tool = GetSystemInfoTool::new();
108        assert_eq!(tool.definition().name, "get_system_info");
109        assert!(!tool.definition().dangerous);
110    }
111}