1use std::{collections::BTreeMap, sync::Arc};
2
3use async_trait::async_trait;
4use serde::Serialize;
5
6pub 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}