Skip to main content

lash_core/plugin/session_obj/
tools.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::Arc;
3
4use super::*;
5
6fn merge_string_array(
7    obj: &mut serde_json::Map<String, serde_json::Value>,
8    key: &str,
9    values: Vec<String>,
10) {
11    let mut existing = obj
12        .remove(key)
13        .and_then(|value| value.as_array().cloned())
14        .unwrap_or_default()
15        .into_iter()
16        .filter_map(|value| value.as_str().map(str::to_string))
17        .collect::<BTreeSet<_>>();
18    existing.extend(
19        values
20            .into_iter()
21            .map(|value| value.trim().to_string())
22            .filter(|value| !value.is_empty()),
23    );
24    if !existing.is_empty() {
25        obj.insert(key.to_string(), serde_json::json!(existing));
26    }
27}
28
29fn apply_tool_discovery_contributions(
30    catalog: &mut [serde_json::Value],
31    contributions: impl IntoIterator<Item = ToolDiscoveryContribution>,
32) {
33    let mut by_name = BTreeMap::new();
34    for (idx, tool) in catalog.iter().enumerate() {
35        if let Some(name) = tool.get("name").and_then(serde_json::Value::as_str) {
36            by_name.insert(name.to_string(), idx);
37        }
38    }
39
40    for contribution in contributions {
41        for patch in contribution.tools {
42            let Some(idx) = by_name.get(&patch.tool_name).copied() else {
43                continue;
44            };
45            let Some(obj) = catalog[idx].as_object_mut() else {
46                continue;
47            };
48            if let Some(namespace) = patch
49                .namespace
50                .map(|value| value.trim().to_string())
51                .filter(|value| !value.is_empty())
52            {
53                obj.insert("namespace".to_string(), serde_json::json!(namespace));
54            }
55            merge_string_array(obj, "aliases", patch.aliases);
56        }
57    }
58}
59
60impl PluginSession {
61    pub fn resolved_tool_catalog(
62        &self,
63        session_id: &str,
64    ) -> Result<Arc<crate::ToolCatalog>, PluginError> {
65        let tools = self.tools.tool_manifests();
66        let contract_provider = Arc::clone(&self.tools);
67        let resolve_contract: lash_sansio::ToolContractResolver =
68            Arc::new(move |name: &str| contract_provider.resolve_contract(name));
69        Ok(Arc::new(self.resolve_tool_catalog(ToolCatalogContext {
70            session_id: session_id.to_string(),
71            tools,
72            resolve_contract: Some(Arc::clone(&resolve_contract)),
73            tool_access: self.tool_access.clone(),
74            subagent: self.subagent.clone(),
75            lashlang_abilities: self.lashlang_abilities,
76        })?))
77    }
78
79    pub fn tool_catalog(&self, session_id: &str) -> Result<Vec<serde_json::Value>, PluginError> {
80        let catalog = self.resolved_tool_catalog(session_id)?;
81        let mut catalog =
82            crate::tool_registry::project_tool_catalog(catalog.searchable_tools_iter().cloned());
83        let contributions = collect_owned_sync(
84            &self.contributions.tool_discovery_contributors,
85            ToolDiscoveryContext {
86                session_id: session_id.to_string(),
87                catalog: catalog.clone(),
88            },
89            |hook, ctx| hook(ctx),
90        )
91        .unwrap_or_else(|err| {
92            tracing::warn!("failed to resolve tool discovery metadata: {err}");
93            Vec::new()
94        });
95        apply_tool_discovery_contributions(
96            &mut catalog,
97            contributions.into_iter().map(|owned| owned.value),
98        );
99        Ok(catalog)
100    }
101
102    pub fn resolve_tool_catalog(
103        &self,
104        ctx: ToolCatalogContext,
105    ) -> Result<crate::ToolCatalog, PluginError> {
106        let mut contributions = collect_owned_sync(
107            &self.contributions.tool_catalog_contributors,
108            ToolCatalogContext {
109                session_id: ctx.session_id.clone(),
110                tools: ctx.tools.clone(),
111                resolve_contract: ctx.resolve_contract.clone(),
112                tool_access: ctx.tool_access.clone(),
113                subagent: ctx.subagent.clone(),
114                lashlang_abilities: ctx.lashlang_abilities,
115            },
116            |hook, ctx| hook(ctx),
117        )?
118        .into_iter()
119        .map(|owned| owned.value)
120        .collect::<Vec<_>>();
121        contributions.push(self.tool_catalog_overlay.clone());
122        let (tools, resolve_contract) = if ctx.tool_access.tools.is_empty() {
123            (ctx.tools, ctx.resolve_contract)
124        } else {
125            let contracts = ctx
126                .tool_access
127                .tools
128                .iter()
129                .map(|tool| (tool.name().to_string(), Arc::new(tool.contract())))
130                .collect::<BTreeMap<_, _>>();
131            (
132                ctx.tool_access
133                    .tools
134                    .iter()
135                    .map(|tool| tool.manifest())
136                    .collect(),
137                Some(Arc::new(move |name: &str| contracts.get(name).cloned())
138                    as lash_sansio::ToolContractResolver),
139            )
140        };
141        let authority_hidden_tools = tools
142            .iter()
143            .filter(|tool| ctx.tool_access.hides(&tool.name))
144            .map(|tool| tool.name.clone())
145            .collect::<BTreeSet<_>>();
146        if !authority_hidden_tools.is_empty() {
147            contributions.push(ToolCatalogContribution {
148                overrides: authority_hidden_tools
149                    .into_iter()
150                    .map(|tool_name| ToolCatalogOverride {
151                        tool_name,
152                        availability: Some(crate::ToolAvailability::Off),
153                    })
154                    .collect(),
155                ..Default::default()
156            });
157        }
158        Ok(crate::build_tool_catalog(crate::ToolCatalogBuildInput {
159            tools,
160            resolve_contract,
161            contributions,
162        }))
163    }
164}