swink_agent/
sub_agent.rs

1//! Sub-agent tool wrapper for multi-agent composition.
2//!
3//! [`SubAgent`] implements [`AgentTool`], allowing an agent to be used as a tool
4//! within a parent agent. On each `execute()` call, it constructs a fresh
5//! [`Agent`] from a factory closure, runs it, and maps the result.
6
7use std::sync::Arc;
8
9use serde_json::{Value, json};
10use tokio_util::sync::CancellationToken;
11
12use crate::agent::{Agent, AgentOptions};
13use crate::stream::StreamFn;
14use crate::tool::{AgentTool, AgentToolResult, ToolFuture};
15use crate::types::{AgentResult, ContentBlock, ModelSpec, StopReason};
16
17// ─── Type aliases ───────────────────────────────────────────────────────────
18
19type OptionsFactoryFn = Arc<dyn Fn() -> AgentOptions + Send + Sync>;
20type MapResultFn = Arc<dyn Fn(AgentResult) -> AgentToolResult + Send + Sync>;
21
22// ─── SubAgent ───────────────────────────────────────────────────────────────
23
24/// A tool that wraps an agent, enabling multi-agent composition.
25///
26/// When executed, constructs a fresh [`Agent`] via the `options_factory`,
27/// sends the prompt extracted from tool call params, and maps the
28/// [`AgentResult`] into an [`AgentToolResult`].
29pub struct SubAgent {
30    name: String,
31    label: String,
32    description: String,
33    schema: Value,
34    requires_approval: bool,
35    options_factory: Option<OptionsFactoryFn>,
36    map_result: MapResultFn,
37}
38
39impl SubAgent {
40    /// Start building a sub-agent tool with the given identity.
41    ///
42    /// Defaults to a schema that accepts a `prompt` string parameter.
43    /// Use [`with_options`](Self::with_options) to configure the inner agent.
44    #[must_use]
45    pub fn new(
46        name: impl Into<String>,
47        label: impl Into<String>,
48        description: impl Into<String>,
49    ) -> Self {
50        Self {
51            name: name.into(),
52            label: label.into(),
53            description: description.into(),
54            schema: json!({
55                "type": "object",
56                "properties": {
57                    "prompt": {
58                        "type": "string",
59                        "description": "The prompt to send to the sub-agent"
60                    }
61                },
62                "required": ["prompt"]
63            }),
64            requires_approval: false,
65            options_factory: None,
66            map_result: Arc::new(default_map_result),
67        }
68    }
69
70    /// Convenience constructor that builds a fully configured sub-agent.
71    ///
72    /// Creates an `AgentOptions::new_simple()` internally with the provided
73    /// system prompt, model, and stream function.
74    #[must_use]
75    pub fn simple(
76        name: impl Into<String>,
77        label: impl Into<String>,
78        description: impl Into<String>,
79        system_prompt: impl Into<String>,
80        model: ModelSpec,
81        stream_fn: Arc<dyn StreamFn>,
82    ) -> Self {
83        let system_prompt = system_prompt.into();
84        Self::new(name, label, description).with_options(move || {
85            AgentOptions::new_simple(system_prompt.clone(), model.clone(), Arc::clone(&stream_fn))
86        })
87    }
88
89    /// Set a custom JSON Schema for the tool parameters.
90    #[must_use]
91    pub fn with_schema(mut self, schema: Value) -> Self {
92        self.schema = schema;
93        self
94    }
95
96    /// Set whether this tool requires approval before execution.
97    #[must_use]
98    pub const fn with_requires_approval(mut self, requires: bool) -> Self {
99        self.requires_approval = requires;
100        self
101    }
102
103    /// Set the factory closure that creates agent options for each execution.
104    #[must_use]
105    pub fn with_options(mut self, f: impl Fn() -> AgentOptions + Send + Sync + 'static) -> Self {
106        self.options_factory = Some(Arc::new(f));
107        self
108    }
109
110    /// Set a custom result mapper from [`AgentResult`] to [`AgentToolResult`].
111    #[must_use]
112    pub fn with_map_result(
113        mut self,
114        f: impl Fn(AgentResult) -> AgentToolResult + Send + Sync + 'static,
115    ) -> Self {
116        self.map_result = Arc::new(f);
117        self
118    }
119}
120
121/// Default result mapper: extracts text from the last assistant message.
122fn default_map_result(result: AgentResult) -> AgentToolResult {
123    if result.stop_reason == StopReason::Error {
124        let error_text = result
125            .error
126            .unwrap_or_else(|| "sub-agent ended with error".to_owned());
127        return AgentToolResult::error(error_text);
128    }
129
130    // Extract text from all messages (last assistant message will have the answer)
131    let text = result
132        .messages
133        .iter()
134        .rev()
135        .find_map(|msg| {
136            if let crate::types::AgentMessage::Llm(crate::types::LlmMessage::Assistant(a)) = msg {
137                let t = ContentBlock::extract_text(&a.content);
138                if t.is_empty() { None } else { Some(t) }
139            } else {
140                None
141            }
142        })
143        .unwrap_or_else(|| "sub-agent produced no text output".to_owned());
144
145    AgentToolResult::text(text)
146}
147
148impl AgentTool for SubAgent {
149    fn name(&self) -> &str {
150        &self.name
151    }
152
153    fn label(&self) -> &str {
154        &self.label
155    }
156
157    fn description(&self) -> &str {
158        &self.description
159    }
160
161    fn parameters_schema(&self) -> &Value {
162        &self.schema
163    }
164
165    fn requires_approval(&self) -> bool {
166        self.requires_approval
167    }
168
169    fn execute(
170        &self,
171        _tool_call_id: &str,
172        params: Value,
173        cancellation_token: CancellationToken,
174        _on_update: Option<Box<dyn Fn(AgentToolResult) + Send + Sync>>,
175        _state: std::sync::Arc<std::sync::RwLock<crate::SessionState>>,
176        _credential: Option<crate::credential::ResolvedCredential>,
177    ) -> ToolFuture<'_> {
178        let options_factory = self.options_factory.clone();
179        let map_result = Arc::clone(&self.map_result);
180        Box::pin(async move {
181            let Some(options_factory) = options_factory else {
182                return AgentToolResult::error(
183                    "Sub-agent options were not configured; call with_options() or simple().",
184                );
185            };
186
187            let options = options_factory();
188            let mut agent = Agent::new(options);
189            let prompt = params["prompt"].as_str().unwrap_or("").to_owned();
190            let result = tokio::select! {
191                r = agent.prompt_text(prompt) => r,
192                () = cancellation_token.cancelled() => {
193                    agent.abort();
194                    return AgentToolResult::error("Sub-agent cancelled.");
195                }
196            };
197            match result {
198                Ok(r) => map_result(r),
199                Err(e) => AgentToolResult::error(format!("Sub-agent error: {e}")),
200            }
201        })
202    }
203}
204
205impl std::fmt::Debug for SubAgent {
206    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207        f.debug_struct("SubAgent")
208            .field("name", &self.name)
209            .field("label", &self.label)
210            .field("description", &self.description)
211            .finish_non_exhaustive()
212    }
213}
214
215// ─── Compile-time Send + Sync assertion ─────────────────────────────────────
216
217const _: () = {
218    const fn assert_send_sync<T: Send + Sync>() {}
219    assert_send_sync::<SubAgent>();
220};
swink_agent/sub_agent.rs

swink_agent/
sub_agent.rs
