Skip to main content

atomr_agents_tool/
tool_return.rs

1//! `ToolReturn` — richer return type for tools that need to drive
2//! graph control flow or separate model-visible content from large
3//! artifacts.
4//!
5//! The base `Tool` trait's `invoke` returns `Value` for backwards
6//! compatibility. Tools that want richer behavior implement
7//! `RichTool::invoke_rich` and the agent loop will pick it up via
8//! the registry-side adapter.
9
10use async_trait::async_trait;
11use atomr_agents_core::{InvokeCtx, Result, Value};
12use serde::{Deserialize, Serialize};
13
14use crate::descriptor::ToolDescriptor;
15use crate::r#trait::Tool;
16
17/// What a richer tool returns. Agents map this back into the message
18/// sequence and (optionally) the workflow state.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub enum ToolReturn {
21    /// Plain content — model sees it as a `Role::Tool` message.
22    Content(Value),
23    /// Both model-visible content and an out-of-band artifact (e.g. a
24    /// large blob). The artifact is stashed by the runner under a
25    /// tool-named slot for later retrieval; only `content` enters the
26    /// next prompt turn.
27    ContentAndArtifact { content: Value, artifact: Value },
28    /// Drive the surrounding harness/graph: send a control instruction.
29    Command(ToolControl),
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub enum ToolControl {
34    /// Hand off control to another agent / handler by id.
35    Handoff { target: String, payload: Value },
36    /// Terminate the current turn early with this value.
37    Done(Value),
38    /// Update one or more workflow channels.
39    Update(Vec<(String, Value)>),
40}
41
42#[async_trait]
43pub trait RichTool: Send + Sync + 'static {
44    fn descriptor(&self) -> &ToolDescriptor;
45    async fn invoke_rich(&self, args: Value, ctx: &InvokeCtx) -> Result<ToolReturn>;
46}
47
48/// Any `RichTool` is automatically a `Tool`: `invoke` projects
49/// `ToolReturn::Content` (or the `content` field of
50/// `ContentAndArtifact`); other variants surface as a synthetic
51/// content carrying the control payload, so legacy callers keep
52/// working.
53#[async_trait]
54impl<T: RichTool> Tool for T {
55    fn descriptor(&self) -> &ToolDescriptor {
56        RichTool::descriptor(self)
57    }
58
59    async fn invoke(&self, args: Value, ctx: &InvokeCtx) -> Result<Value> {
60        match RichTool::invoke_rich(self, args, ctx).await? {
61            ToolReturn::Content(v) => Ok(v),
62            ToolReturn::ContentAndArtifact { content, .. } => Ok(content),
63            ToolReturn::Command(c) => Ok(serde_json::to_value(c).unwrap_or(Value::Null)),
64        }
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71    use crate::descriptor::ToolSchema;
72    use atomr_agents_core::{
73        CallCtx, InvokeCtx, IterationBudget, MoneyBudget, TimeBudget, TokenBudget, ToolId,
74    };
75    use std::time::Duration;
76
77    struct Handoff {
78        d: ToolDescriptor,
79    }
80
81    #[async_trait]
82    impl RichTool for Handoff {
83        fn descriptor(&self) -> &ToolDescriptor {
84            &self.d
85        }
86        async fn invoke_rich(&self, _args: Value, _ctx: &InvokeCtx) -> Result<ToolReturn> {
87            Ok(ToolReturn::Command(ToolControl::Handoff {
88                target: "specialist".into(),
89                payload: serde_json::json!({"why": "complex"}),
90            }))
91        }
92    }
93
94    fn ictx() -> InvokeCtx {
95        InvokeCtx {
96            call: CallCtx {
97                agent_id: None,
98                tokens: TokenBudget::new(1000),
99                time: TimeBudget::new(Duration::from_secs(5)),
100                money: MoneyBudget::from_usd(0.10),
101                iterations: IterationBudget::new(5),
102                trace: vec![],
103                extensions: Default::default(),
104            },
105            tool_call_id: "t1".into(),
106            raw_args: Value::Null,
107        }
108    }
109
110    #[tokio::test]
111    async fn rich_tool_acts_as_plain_tool() {
112        let t = Handoff {
113            d: ToolDescriptor {
114                id: ToolId::from("handoff"),
115                name: "handoff".into(),
116                description: "delegate to specialist".into(),
117                schema: ToolSchema::empty_object(),
118            },
119        };
120        let v = Tool::invoke(&t, Value::Null, &ictx()).await.unwrap();
121        assert!(v.is_object()); // serialized ToolControl
122    }
123
124    #[tokio::test]
125    async fn rich_invoke_returns_command() {
126        let t = Handoff {
127            d: ToolDescriptor {
128                id: ToolId::from("handoff"),
129                name: "handoff".into(),
130                description: "delegate".into(),
131                schema: ToolSchema::empty_object(),
132            },
133        };
134        let r = t.invoke_rich(Value::Null, &ictx()).await.unwrap();
135        assert!(matches!(r, ToolReturn::Command(ToolControl::Handoff { .. })));
136    }
137}