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 tool_surface(&self, session_id: &str) -> Result<Arc<crate::ToolSurface>, PluginError> {
62        let tools = self.tools.tool_manifests();
63        let contract_provider = Arc::clone(&self.tools);
64        let resolve_contract: lash_sansio::ToolContractResolver =
65            Arc::new(move |name: &str| contract_provider.resolve_contract(name));
66        Ok(Arc::new(self.resolve_tool_surface(ToolSurfaceContext {
67            session_id: session_id.to_string(),
68            tools,
69            resolve_contract: Some(Arc::clone(&resolve_contract)),
70            tool_access: self.tool_access.clone(),
71            subagent: self.subagent.clone(),
72            lashlang_abilities: self.lashlang_abilities,
73        })?))
74    }
75
76    pub fn tool_catalog(&self, session_id: &str) -> Result<Vec<serde_json::Value>, PluginError> {
77        let surface = self.tool_surface(session_id)?;
78        let mut catalog =
79            crate::tool_registry::project_tool_catalog(surface.searchable_tools_iter().cloned());
80        let contributions = collect_owned_sync(
81            &self.contributions.tool_discovery_contributors,
82            ToolDiscoveryContext {
83                session_id: session_id.to_string(),
84                catalog: catalog.clone(),
85            },
86            |hook, ctx| hook(ctx),
87        )
88        .unwrap_or_else(|err| {
89            tracing::warn!("failed to resolve tool discovery metadata: {err}");
90            Vec::new()
91        });
92        apply_tool_discovery_contributions(
93            &mut catalog,
94            contributions.into_iter().map(|owned| owned.value),
95        );
96        Ok(catalog)
97    }
98
99    pub fn resolve_tool_surface(
100        &self,
101        ctx: ToolSurfaceContext,
102    ) -> Result<crate::ToolSurface, PluginError> {
103        let mut contributions = collect_owned_sync(
104            &self.contributions.tool_surface_contributors,
105            ToolSurfaceContext {
106                session_id: ctx.session_id.clone(),
107                tools: ctx.tools.clone(),
108                resolve_contract: ctx.resolve_contract.clone(),
109                tool_access: ctx.tool_access.clone(),
110                subagent: ctx.subagent.clone(),
111                lashlang_abilities: ctx.lashlang_abilities,
112            },
113            |hook, ctx| hook(ctx),
114        )?
115        .into_iter()
116        .map(|owned| owned.value)
117        .collect::<Vec<_>>();
118        contributions.push(self.tool_surface_overlay.clone());
119        let (tools, resolve_contract) = if ctx.tool_access.tools.is_empty() {
120            (ctx.tools, ctx.resolve_contract)
121        } else {
122            let contracts = ctx
123                .tool_access
124                .tools
125                .iter()
126                .map(|tool| (tool.name().to_string(), Arc::new(tool.contract())))
127                .collect::<BTreeMap<_, _>>();
128            (
129                ctx.tool_access
130                    .tools
131                    .iter()
132                    .map(|tool| tool.manifest())
133                    .collect(),
134                Some(Arc::new(move |name: &str| contracts.get(name).cloned())
135                    as lash_sansio::ToolContractResolver),
136            )
137        };
138        let authority_hidden_tools = tools
139            .iter()
140            .filter(|tool| ctx.tool_access.hides(&tool.name))
141            .map(|tool| tool.name.clone())
142            .collect::<BTreeSet<_>>();
143        if !authority_hidden_tools.is_empty() {
144            contributions.push(ToolSurfaceContribution {
145                overrides: authority_hidden_tools
146                    .into_iter()
147                    .map(|tool_name| ToolSurfaceOverride {
148                        tool_name,
149                        availability: Some(crate::ToolAvailability::Off),
150                    })
151                    .collect(),
152                ..Default::default()
153            });
154        }
155        Ok(crate::build_tool_surface(crate::ToolSurfaceBuildInput {
156            tools,
157            resolve_contract,
158            contributions,
159        }))
160    }
161}