oxios_kernel/kernel_handle/
agent_api.rs

1//! Agent API — agent lifecycle, budget, memory.
2
3use crate::budget::{BudgetExceeded, BudgetInfo, BudgetLimit, BudgetManager};
4use crate::memory::store::{HnswMemoryIndex, SemanticHit};
5use crate::memory::{MemoryEntry, MemoryManager, MemoryType};
6use crate::supervisor::Supervisor;
7use crate::types::AgentId;
8use std::sync::Arc;
9
10/// Agent management system calls.
11pub struct AgentApi {
12    pub(crate) supervisor: Arc<dyn Supervisor>,
13    pub(crate) budget_manager: Arc<BudgetManager>,
14    pub(crate) memory_manager: Arc<MemoryManager>,
15    /// HNSW index for semantic search (optional, initialized on demand).
16    pub(crate) hnsw_index: Option<Arc<HnswMemoryIndex>>,
17}
18
19impl AgentApi {
20    /// Create a new AgentApi.
21    pub fn new(
22        supervisor: Arc<dyn Supervisor>,
23        budget_manager: Arc<BudgetManager>,
24        memory_manager: Arc<MemoryManager>,
25    ) -> Self {
26        Self {
27            supervisor,
28            budget_manager,
29            memory_manager,
30            hnsw_index: None,
31        }
32    }
33
34    /// Attach an HNSW index for semantic search.
35    pub fn set_hnsw_index(&mut self, index: Arc<HnswMemoryIndex>) {
36        self.hnsw_index = Some(index);
37    }
38    /// List running agents.
39    pub async fn list(&self) -> anyhow::Result<Vec<crate::types::AgentInfo>> {
40        self.supervisor
41            .list()
42            .await
43            .map_err(|e| anyhow::anyhow!("supervisor: {e}"))
44    }
45
46    /// Kill a running agent.
47    pub async fn kill(&self, agent_id: &str) -> anyhow::Result<()> {
48        let id = uuid::Uuid::parse_str(agent_id)
49            .map_err(|e| anyhow::anyhow!("invalid agent id: {e}"))?;
50        self.supervisor
51            .kill(id)
52            .await
53            .map_err(|e| anyhow::anyhow!("supervisor: {e}"))
54    }
55
56    /// Check budget for an agent.
57    pub fn check_budget(&self, agent_id: &AgentId) -> BudgetInfo {
58        self.budget_manager.remaining(agent_id)
59    }
60
61    /// Set budget for an agent.
62    pub fn set_budget(&self, limit: BudgetLimit) {
63        self.budget_manager.set_budget(limit);
64    }
65
66    /// Remove budget for an agent.
67    pub fn remove_budget(&self, agent_id: &AgentId) {
68        self.budget_manager.remove_budget(agent_id);
69    }
70
71    /// Reserve tokens for an agent.
72    pub fn reserve_budget(&self, agent_id: &AgentId, tokens: u64) -> Result<(), BudgetExceeded> {
73        self.budget_manager.reserve(agent_id, tokens)
74    }
75
76    /// Reset budget window for an agent.
77    pub fn reset_budget(&self, agent_id: &AgentId) {
78        self.budget_manager.reset_window(agent_id);
79    }
80
81    /// Get memory stats.
82    pub async fn memory_stats(&self) -> (usize, usize) {
83        (
84            self.memory_manager.vector_index_size(),
85            self.memory_manager.total_entries().await,
86        )
87    }
88
89    /// Store a memory entry.
90    pub async fn remember(&self, entry: MemoryEntry) -> anyhow::Result<String> {
91        self.memory_manager.remember(entry).await
92    }
93
94    /// Search memory entries.
95    pub async fn search_memory(
96        &self,
97        query: &str,
98        memory_type: Option<MemoryType>,
99        limit: usize,
100    ) -> anyhow::Result<Vec<MemoryEntry>> {
101        self.memory_manager.search(query, memory_type, limit).await
102    }
103
104    /// Semantic search using HNSW index.
105    ///
106    /// Falls back to regular search if HNSW index is not available.
107    pub async fn semantic_search_memory(
108        &self,
109        query: &str,
110        memory_type: Option<MemoryType>,
111        limit: usize,
112    ) -> anyhow::Result<Vec<SemanticHit>> {
113        if let Some(ref hnsw) = self.hnsw_index {
114            self.memory_manager
115                .semantic_search(query, memory_type, limit, hnsw)
116                .await
117        } else {
118            // Fallback to regular search, wrap results
119            let entries = self.search_memory(query, memory_type, limit).await?;
120            Ok(entries
121                .into_iter()
122                .map(|entry| SemanticHit {
123                    entry,
124                    distance: 0.0,
125                    similarity: 0.0,
126                })
127                .collect())
128        }
129    }
130
131    /// Memory manager reference.
132    pub fn memory_manager(&self) -> &Arc<MemoryManager> {
133        &self.memory_manager
134    }
135
136    /// Rebuild the HNSW index from all stored memories.
137    pub async fn rebuild_hnsw_index(&self) -> anyhow::Result<usize> {
138        if let Some(ref hnsw) = self.hnsw_index {
139            self.memory_manager.rebuild_hnsw_index(hnsw).await
140        } else {
141            Err(anyhow::anyhow!("HNSW index not initialized"))
142        }
143    }
144}
oxios_kernel/kernel_handle/agent_api.rs

oxios_kernel/kernel_handle/
agent_api.rs
