oxios_kernel/tools/builtin/
budget_tool.rs

1//! Budget tool — wraps `AgentApi` budget methods behind the `AgentTool` interface.
2//!
3//! Provides agents with budget management capabilities.
4//! Actions: check, set, reserve, reset.
5//!
6//! ## Example
7//!
8//! ```json
9//! { "action": "check", "agent_id": "uuid" }
10//! { "action": "set", "agent_id": "uuid", "limit": 100000 }
11//! { "action": "reserve", "agent_id": "uuid", "tokens": 500 }
12//! { "action": "reset", "agent_id": "uuid" }
13//! ```
14
15use std::sync::Arc;
16
17use async_trait::async_trait;
18use oxi_sdk::{AgentTool, AgentToolResult, ToolContext};
19use serde_json::{json, Value};
20use tokio::sync::oneshot;
21
22use crate::budget::{BudgetLimit, BudgetManager};
23use crate::kernel_handle::KernelHandle;
24use crate::types::AgentId;
25
26/// Agent tool for budget management.
27///
28/// Wraps the budget-related methods of the `AgentApi` domain. Allows agents
29/// to check, configure, and manage token/call budgets for agents.
30///
31/// ## Actions
32///
33/// | Action    | Description                  | Required params | Optional params |
34/// |-----------|------------------------------|-----------------|-----------------|
35/// | `check`   | Check remaining budget       | `agent_id`      | —               |
36/// | `set`     | Set budget limit for agent   | `agent_id`      | `limit`         |
37/// | `reserve` | Reserve tokens from budget   | `agent_id`      | `tokens`        |
38/// | `reset`   | Reset budget window          | `agent_id`      | —               |
39pub struct BudgetTool {
40    budget_manager: Arc<BudgetManager>,
41}
42
43impl BudgetTool {
44    /// Create a new `BudgetTool` from a `KernelHandle`.
45    ///
46    /// Extracts the `BudgetManager` Arc from the kernel's Agent API.
47    pub fn from_kernel(kernel: &KernelHandle) -> Self {
48        Self {
49            budget_manager: kernel.agents.budget_manager.clone(),
50        }
51    }
52}
53
54impl std::fmt::Debug for BudgetTool {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.debug_struct("BudgetTool").finish()
57    }
58}
59
60#[async_trait]
61impl AgentTool for BudgetTool {
62    fn name(&self) -> &str {
63        "budget"
64    }
65
66    fn label(&self) -> &str {
67        "Budget"
68    }
69
70    fn description(&self) -> &'static str {
71        "Manage agent budgets — check remaining tokens/calls, set limits, reserve tokens, or reset the budget window. \
72         Actions: check, set, reserve, reset."
73    }
74
75    fn parameters_schema(&self) -> Value {
76        json!({
77            "type": "object",
78            "properties": {
79                "action": {
80                    "type": "string",
81                    "enum": ["check", "set", "reserve", "reset"],
82                    "description": "Budget operation to perform"
83                },
84                "agent_id": {
85                    "type": "string",
86                    "description": "Agent UUID to operate on"
87                },
88                "limit": {
89                    "type": "integer",
90                    "description": "Token budget limit (set action only, default: 100000)"
91                },
92                "tokens": {
93                    "type": "integer",
94                    "description": "Number of tokens to reserve (reserve action only)"
95                }
96            },
97            "required": ["action", "agent_id"]
98        })
99    }
100
101    async fn execute(
102        &self,
103        _tool_call_id: &str,
104        params: Value,
105        _signal: Option<oneshot::Receiver<()>>,
106        _ctx: &ToolContext,
107    ) -> Result<AgentToolResult, String> {
108        let action = params
109            .get("action")
110            .and_then(|v| v.as_str())
111            .ok_or_else(|| "Missing required parameter: action".to_string())?;
112
113        let agent_id_str = params
114            .get("agent_id")
115            .and_then(|v| v.as_str())
116            .ok_or_else(|| "Missing required parameter: agent_id".to_string())?;
117
118        let agent_id: AgentId = match uuid::Uuid::parse_str(agent_id_str) {
119            Ok(id) => id,
120            Err(e) => return Ok(AgentToolResult::error(format!("Invalid agent_id: {e}"))),
121        };
122
123        match action {
124            "check" => {
125                let info = self.budget_manager.remaining(&agent_id);
126                Ok(AgentToolResult::success(
127                    serde_json::to_string_pretty(&json!({
128                        "agent_id": agent_id_str,
129                        "tokens_remaining": info.tokens_remaining,
130                        "calls_remaining": info.calls_remaining,
131                        "window_remaining_secs": info.window_remaining_secs,
132                        "is_exhausted": info.is_exhausted,
133                    }))
134                    .unwrap_or_default(),
135                ))
136            }
137
138            "set" => {
139                let token_limit = params["limit"].as_u64().unwrap_or(100_000);
140
141                let limit = BudgetLimit {
142                    agent_id,
143                    token_budget: token_limit,
144                    calls_budget: 1_000,
145                    window_secs: 3_600, // 1 hour default
146                };
147
148                self.budget_manager.set_budget(limit);
149
150                Ok(AgentToolResult::success(format!(
151                    "Budget set for agent '{agent_id_str}': {token_limit} tokens, 1000 calls, 1h window.",
152                )))
153            }
154
155            "reserve" => {
156                let tokens = params["tokens"].as_u64().unwrap_or(0);
157                if tokens == 0 {
158                    return Ok(AgentToolResult::error(
159                        "reserve requires 'tokens' parameter (> 0)",
160                    ));
161                }
162
163                match self.budget_manager.reserve(&agent_id, tokens) {
164                    Ok(()) => Ok(AgentToolResult::success(format!(
165                        "Reserved {tokens} tokens for agent '{agent_id_str}'.",
166                    ))),
167                    Err(exceeded) => Ok(AgentToolResult::error(format!(
168                        "Budget exceeded: {exceeded}"
169                    ))),
170                }
171            }
172
173            "reset" => {
174                self.budget_manager.reset_window(&agent_id);
175                Ok(AgentToolResult::success(format!(
176                    "Budget window reset for agent '{agent_id_str}'.",
177                )))
178            }
179
180            other => Err(format!(
181                "Unknown budget action '{other}'. Valid: check, set, reserve, reset"
182            )),
183        }
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    #[test]
192    fn test_schema_structure() {
193        let schema = json!({
194            "type": "object",
195            "properties": {
196                "action": {
197                    "type": "string",
198                    "enum": ["check", "set", "reserve", "reset"]
199                },
200                "agent_id": { "type": "string" },
201                "limit": { "type": "integer" },
202                "tokens": { "type": "integer" }
203            },
204            "required": ["action", "agent_id"]
205        });
206
207        let actions = schema["properties"]["action"]["enum"].as_array().unwrap();
208        assert_eq!(actions.len(), 4);
209        assert!(actions.iter().any(|a| a == "check"));
210        assert!(actions.iter().any(|a| a == "set"));
211        assert!(actions.iter().any(|a| a == "reserve"));
212        assert!(actions.iter().any(|a| a == "reset"));
213
214        let required = schema["required"].as_array().unwrap();
215        assert!(required.iter().any(|r| r.as_str() == Some("action")));
216        assert!(required.iter().any(|r| r.as_str() == Some("agent_id")));
217    }
218}
oxios_kernel/tools/builtin/budget_tool.rs

oxios_kernel/tools/builtin/
budget_tool.rs
