oino_core/session/
nested.rs

1//! Nested tool calls (`ctx.executeTool`, `nested-tool-calls.ts`): a tool runs another
2//! callable tool through the full pipeline (validation, `tool_call`/`tool_result` hooks) with
3//! id `<parent>/<n>`; calls and usage are recorded on the parent's result message.
4
5use std::collections::HashMap;
6use std::time::Instant;
7
8use oino_agent::{
9    AfterToolCall, AfterToolCallContext, AfterToolCallResult, AgentContext, BeforeToolCall,
10    BeforeToolCallContext, BeforeToolCallResult, HookError, RunToolCallOptions, ToolCallOutcome,
11    run_tool_call,
12};
13use oino_ai::{
14    AssistantMessage, JsonObject, NestedStatus, NestedToolCallRecord, NestedToolCalls, ToolCall,
15    Usage, UserContent,
16};
17use tokio_util::sync::CancellationToken;
18
19use super::{AgentSession, lock};
20use crate::events::AgentSessionEvent;
21
22/// The nested call limits constant.
23const MAX_CALLS: usize = 256;
24const MAX_ARGUMENT_BYTES_PER_CALL: u64 = 8 * 1024;
25const MAX_ARGUMENT_BYTES_TOTAL: u64 = 32 * 1024;
26const MAX_ERROR_CHARS: usize = 500;
27
28/// The nested-call recorder for one top-level call.
29pub(crate) struct Recorder {
30    calls: Vec<(NestedToolCallRecord, Instant)>,
31    complete: bool,
32    argument_bytes: u64,
33    usage: Option<Usage>,
34    next_id: u64,
35}
36
37impl Default for Recorder {
38    fn default() -> Self {
39        Self {
40            calls: Vec::new(),
41            complete: true,
42            argument_bytes: 0,
43            usage: None,
44            next_id: 1,
45        }
46    }
47}
48
49impl Recorder {
50    fn start(&mut self, call: &ToolCall) -> Option<usize> {
51        if self.calls.len() >= MAX_CALLS {
52            self.complete = false;
53            return None;
54        }
55        let json = serde_json::to_string(&call.arguments).unwrap_or_default();
56        let bytes = json.len() as u64;
57        let mut record = NestedToolCallRecord {
58            id: call.id.clone(),
59            name: call.name.clone(),
60            arguments: None,
61            arguments_bytes: None,
62            status: NestedStatus::Unfinished,
63            duration_ms: None,
64            error: None,
65        };
66        if bytes > MAX_ARGUMENT_BYTES_PER_CALL
67            || self.argument_bytes + bytes > MAX_ARGUMENT_BYTES_TOTAL
68        {
69            record.arguments_bytes = Some(bytes);
70            self.complete = false;
71        } else {
72            record.arguments = Some(call.arguments.clone());
73            self.argument_bytes += bytes;
74        }
75        self.calls.push((record, Instant::now()));
76        Some(self.calls.len() - 1)
77    }
78
79    fn finish(&mut self, index: Option<usize>, outcome: &ToolCallOutcome) {
80        let Some((record, started)) = index.and_then(|i| self.calls.get_mut(i)) else {
81            return;
82        };
83        record.status = if outcome.is_error {
84            NestedStatus::Error
85        } else {
86            NestedStatus::Ok
87        };
88        record.duration_ms = Some(u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX));
89        let text: Vec<&str> = outcome
90            .result
91            .content
92            .iter()
93            .filter_map(|block| match block {
94                UserContent::Text(text) => Some(text.text.as_str()),
95                UserContent::Image(_) => None,
96            })
97            .collect();
98        let text = text.join("\n");
99        if outcome.is_error && !text.is_empty() {
100            record.error = Some(text.chars().take(MAX_ERROR_CHARS).collect());
101        }
102        if let Some(usage) = &outcome.result.usage {
103            match &mut self.usage {
104                Some(total) => total.accumulate(usage),
105                None => self.usage = Some(*usage),
106            }
107        }
108    }
109
110    fn snapshot(&self) -> Option<NestedToolCalls> {
111        if self.calls.is_empty() && self.complete {
112            return None;
113        }
114        let calls: Vec<NestedToolCallRecord> = self.calls.iter().map(|(r, _)| r.clone()).collect();
115        let complete = self.complete && calls.iter().all(|c| c.status != NestedStatus::Unfinished);
116        Some(NestedToolCalls { calls, complete })
117    }
118}
119
120/// Recorders by top-level tool-call id, plus nested-id → top-level id.
121#[derive(Default)]
122pub(crate) struct NestedState {
123    recorders: HashMap<String, Recorder>,
124    roots: HashMap<String, String>,
125    /// Records taken at the result's `message_start`, applied again at `message_end`.
126    taken: HashMap<String, (Option<NestedToolCalls>, Option<Usage>)>,
127}
128
129struct NestedHooks {
130    session: AgentSession,
131    parent: String,
132}
133
134impl BeforeToolCall for NestedHooks {
135    fn before<'a>(
136        &'a self,
137        cx: BeforeToolCallContext<'a>,
138        _cancel: &'a CancellationToken,
139    ) -> futures::future::BoxFuture<'a, Result<Option<BeforeToolCallResult>, HookError>> {
140        Box::pin(async move {
141            self.session
142                .before_tool_call(cx.tool_call, cx.args, Some(&self.parent))
143                .await
144        })
145    }
146}
147
148impl AfterToolCall for NestedHooks {
149    fn after<'a>(
150        &'a self,
151        cx: AfterToolCallContext<'a>,
152        _cancel: &'a CancellationToken,
153    ) -> futures::future::BoxFuture<'a, Result<Option<AfterToolCallResult>, HookError>> {
154        Box::pin(async move { Ok(self.session.after_tool_call(&cx, Some(&self.parent)).await) })
155    }
156}
157
158impl AgentSession {
159    /// `ctx.executeTool(name, args)` for the running call `parent`.
160    pub async fn execute_nested_tool(
161        &self,
162        parent: &str,
163        name: &str,
164        args: JsonObject,
165        cancel: CancellationToken,
166    ) -> ToolCallOutcome {
167        let (call, index) = {
168            let mut state = lock(&self.inner.nested);
169            let root = state
170                .roots
171                .get(parent)
172                .cloned()
173                .unwrap_or_else(|| parent.to_owned());
174            let recorder = state.recorders.entry(root.clone()).or_default();
175            let call = ToolCall {
176                id: format!("{parent}/{}", recorder.next_id),
177                name: name.to_owned(),
178                arguments: args,
179                ..ToolCall::default()
180            };
181            recorder.next_id += 1;
182            let index = recorder.start(&call);
183            state.roots.insert(call.id.clone(), root);
184            (call, index)
185        };
186        self.emit(AgentSessionEvent::ToolExecutionStart {
187            tool_call_id: call.id.clone(),
188            tool_name: name.to_owned(),
189            args: call.arguments.clone(),
190            parent_tool_call_id: Some(parent.to_owned()),
191        });
192        let outcome = match self.last_assistant_message() {
193            None => ToolCallOutcome {
194                tool_call: call.clone(),
195                result: oino_agent::AgentToolResult::error("No assistant message issued this call"),
196                is_error: true,
197                duration_ms: None,
198            },
199            Some(assistant) => {
200                self.run_nested(call.clone(), parent, &assistant, cancel)
201                    .await
202            }
203        };
204        {
205            let mut state = lock(&self.inner.nested);
206            let root = state.roots.remove(&call.id);
207            if let Some(recorder) = root.and_then(|r| state.recorders.get_mut(&r)) {
208                recorder.finish(index, &outcome);
209            }
210        }
211        self.emit(AgentSessionEvent::ToolExecutionEnd {
212            tool_call_id: call.id.clone(),
213            tool_name: name.to_owned(),
214            result: outcome.result.clone(),
215            is_error: outcome.is_error,
216            duration_ms: outcome.duration_ms,
217            parent_tool_call_id: Some(parent.to_owned()),
218        });
219        outcome
220    }
221
222    async fn run_nested(
223        &self,
224        call: ToolCall,
225        parent: &str,
226        assistant: &AssistantMessage,
227        cancel: CancellationToken,
228    ) -> ToolCallOutcome {
229        let active = self.active_tool_names().into_iter().collect();
230        let tools = self.callable_tools(&active);
231        let hooks = std::sync::Arc::new(NestedHooks {
232            session: self.clone(),
233            parent: parent.to_owned(),
234        });
235        let context = AgentContext {
236            messages: self.messages(),
237            tools: self.inner.agent.with_state(|s| s.tools.clone()),
238        };
239        run_tool_call(
240            call,
241            RunToolCallOptions {
242                assistant_message: assistant,
243                context: &context,
244                tools: &tools,
245                before_tool_call: Some(std::sync::Arc::clone(&hooks) as _),
246                after_tool_call: Some(hooks as _),
247                cancel,
248                parent_tool_call_id: Some(parent.to_owned()),
249                on_update: None,
250            },
251        )
252        .await
253    }
254
255    /// Adds the nested calls and usage of a finished top-level call to its result message
256    /// (`takeRecord` at the result's `message_start`; the same record again at `message_end`).
257    pub(crate) fn attach_nested_record(
258        &self,
259        message: &oino_ai::ToolResultMessage,
260        end: bool,
261    ) -> Option<oino_ai::ToolResultMessage> {
262        let mut state = lock(&self.inner.nested);
263        let record = if end {
264            state.taken.remove(&message.tool_call_id)?
265        } else {
266            let recorder = state.recorders.remove(&message.tool_call_id)?;
267            let record = (recorder.snapshot(), recorder.usage);
268            state
269                .taken
270                .insert(message.tool_call_id.clone(), record.clone());
271            record
272        };
273        let mut message = message.clone();
274        if let Some(calls) = record.0 {
275            message.nested_calls = Some(calls);
276        }
277        if let Some(usage) = record.1 {
278            match &mut message.usage {
279                Some(total) => total.accumulate(&usage),
280                None => message.usage = Some(usage),
281            }
282        }
283        Some(message)
284    }
285
286    /// Clears nested-call state at `agent_end`.
287    pub(crate) fn clear_nested(&self) {
288        let mut state = lock(&self.inner.nested);
289        state.recorders.clear();
290        state.roots.clear();
291        state.taken.clear();
292    }
293}
oino_core/session/nested.rs

oino_core/session/
nested.rs
