Skip to main content

loong_kernel/
tool.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use async_trait::async_trait;
4use serde::Serialize;
5
6// Re-export data types from contracts
7pub use loong_contracts::{
8    ToolCoreOutcome, ToolCoreRequest, ToolExtensionOutcome, ToolExtensionRequest, ToolTier,
9};
10
11use crate::errors::ToolPlaneError;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
14#[serde(rename_all = "snake_case")]
15pub enum ToolConcurrencyClass {
16    ReadOnly,
17    Mutating,
18    Unknown,
19}
20
21impl ToolConcurrencyClass {
22    pub const fn as_str(self) -> &'static str {
23        match self {
24            Self::ReadOnly => "read_only",
25            Self::Mutating => "mutating",
26            Self::Unknown => "unknown",
27        }
28    }
29
30    pub const fn requires_serial_execution(self) -> bool {
31        !matches!(self, Self::ReadOnly)
32    }
33}
34
35#[async_trait]
36pub trait CoreToolAdapter: Send + Sync {
37    fn name(&self) -> &str;
38
39    async fn execute_core_tool(
40        &self,
41        request: ToolCoreRequest,
42    ) -> Result<ToolCoreOutcome, ToolPlaneError>;
43}
44
45#[async_trait]
46pub trait ToolExtensionAdapter: Send + Sync {
47    fn name(&self) -> &str;
48
49    async fn execute_tool_extension(
50        &self,
51        request: ToolExtensionRequest,
52        core: &(dyn CoreToolAdapter + Sync),
53    ) -> Result<ToolExtensionOutcome, ToolPlaneError>;
54}
55
56#[derive(Default)]
57pub struct ToolPlane {
58    core_adapters: BTreeMap<String, Arc<dyn CoreToolAdapter>>,
59    extension_adapters: BTreeMap<String, Arc<dyn ToolExtensionAdapter>>,
60    default_core_adapter: Option<String>,
61}
62
63impl ToolPlane {
64    #[must_use]
65    pub fn new() -> Self {
66        Self {
67            core_adapters: BTreeMap::new(),
68            extension_adapters: BTreeMap::new(),
69            default_core_adapter: None,
70        }
71    }
72
73    pub fn register_core_adapter<A: CoreToolAdapter + 'static>(&mut self, adapter: A) {
74        let name = adapter.name().to_owned();
75        if self.default_core_adapter.is_none() {
76            self.default_core_adapter = Some(name.clone());
77        }
78        self.core_adapters.insert(name, Arc::new(adapter));
79    }
80
81    pub fn register_extension_adapter<A: ToolExtensionAdapter + 'static>(&mut self, adapter: A) {
82        let name = adapter.name().to_owned();
83        self.extension_adapters.insert(name, Arc::new(adapter));
84    }
85
86    pub fn set_default_core_adapter(&mut self, name: &str) -> Result<(), ToolPlaneError> {
87        if !self.core_adapters.contains_key(name) {
88            return Err(ToolPlaneError::CoreAdapterNotFound(name.to_owned()));
89        }
90        self.default_core_adapter = Some(name.to_owned());
91        Ok(())
92    }
93
94    #[must_use]
95    pub fn default_core_adapter_name(&self) -> Option<&str> {
96        self.default_core_adapter.as_deref()
97    }
98
99    pub async fn execute_core(
100        &self,
101        core_name: Option<&str>,
102        request: ToolCoreRequest,
103    ) -> Result<ToolCoreOutcome, ToolPlaneError> {
104        let resolved_name = if let Some(name) = core_name {
105            name
106        } else {
107            self.default_core_adapter
108                .as_deref()
109                .ok_or(ToolPlaneError::NoDefaultCoreAdapter)?
110        };
111
112        let adapter = self
113            .core_adapters
114            .get(resolved_name)
115            .ok_or(ToolPlaneError::CoreAdapterNotFound(
116                resolved_name.to_owned(),
117            ))?
118            .clone();
119
120        return adapter.execute_core_tool(request).await;
121    }
122
123    pub async fn execute_extension(
124        &self,
125        extension_name: &str,
126        core_name: Option<&str>,
127        request: ToolExtensionRequest,
128    ) -> Result<ToolExtensionOutcome, ToolPlaneError> {
129        let extension = self
130            .extension_adapters
131            .get(extension_name)
132            .ok_or_else(|| ToolPlaneError::ExtensionNotFound(extension_name.to_owned()))?
133            .clone();
134
135        let resolved_core_name = if let Some(name) = core_name {
136            name
137        } else {
138            self.default_core_adapter
139                .as_deref()
140                .ok_or(ToolPlaneError::NoDefaultCoreAdapter)?
141        };
142
143        let core = self
144            .core_adapters
145            .get(resolved_core_name)
146            .ok_or(ToolPlaneError::CoreAdapterNotFound(
147                resolved_core_name.to_owned(),
148            ))?
149            .clone();
150
151        return extension
152            .execute_tool_extension(request, core.as_ref())
153            .await;
154    }
155}