oxios_kernel/kernel_handle/
agent_api.rs

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

oxios_kernel/kernel_handle/
agent_api.rs
