oxios_kernel/tools/kernel/
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 '{}': {} tokens, 1000 calls, 1h window.",
152                    agent_id_str, token_limit,
153                )))
154            }
155
156            "reserve" => {
157                let tokens = params["tokens"].as_u64().unwrap_or(0);
158                if tokens == 0 {
159                    return Ok(AgentToolResult::error(
160                        "reserve requires 'tokens' parameter (> 0)",
161                    ));
162                }
163
164                match self.budget_manager.reserve(&agent_id, tokens) {
165                    Ok(()) => Ok(AgentToolResult::success(format!(
166                        "Reserved {} tokens for agent '{}'.",
167                        tokens, agent_id_str,
168                    ))),
169                    Err(exceeded) => Ok(AgentToolResult::error(format!(
170                        "Budget exceeded: {}",
171                        exceeded
172                    ))),
173                }
174            }
175
176            "reset" => {
177                self.budget_manager.reset_window(&agent_id);
178                Ok(AgentToolResult::success(format!(
179                    "Budget window reset for agent '{}'.",
180                    agent_id_str,
181                )))
182            }
183
184            other => Err(format!(
185                "Unknown budget action '{}'. Valid: check, set, reserve, reset",
186                other
187            )),
188        }
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    #[test]
197    fn test_schema_structure() {
198        let schema = json!({
199            "type": "object",
200            "properties": {
201                "action": {
202                    "type": "string",
203                    "enum": ["check", "set", "reserve", "reset"]
204                },
205                "agent_id": { "type": "string" },
206                "limit": { "type": "integer" },
207                "tokens": { "type": "integer" }
208            },
209            "required": ["action", "agent_id"]
210        });
211
212        let actions = schema["properties"]["action"]["enum"].as_array().unwrap();
213        assert_eq!(actions.len(), 4);
214        assert!(actions.iter().any(|a| a == "check"));
215        assert!(actions.iter().any(|a| a == "set"));
216        assert!(actions.iter().any(|a| a == "reserve"));
217        assert!(actions.iter().any(|a| a == "reset"));
218
219        let required = schema["required"].as_array().unwrap();
220        assert!(required.iter().any(|r| r.as_str() == Some("action")));
221        assert!(required.iter().any(|r| r.as_str() == Some("agent_id")));
222    }
223}
oxios_kernel/tools/kernel/budget_tool.rs

oxios_kernel/tools/kernel/
budget_tool.rs
