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::store::{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 memory stats.
94    pub async fn memory_stats(&self) -> (usize, usize) {
95        (
96            self.memory_manager.vector_index_size(),
97            self.memory_manager.total_entries().await,
98        )
99    }
100
101    /// Store a memory entry.
102    pub async fn remember(&self, entry: MemoryEntry) -> anyhow::Result<String> {
103        let id = self.memory_manager.remember(entry.clone()).await?;
104
105        // Publish MemoryStored event
106        self.publish(KernelEvent::MemoryStored {
107            id: id.clone(),
108            memory_type: entry.memory_type.label().to_string(),
109            source: entry.source.clone(),
110        });
111
112        Ok(id)
113    }
114
115    /// Search memory entries.
116    pub async fn search_memory(
117        &self,
118        query: &str,
119        memory_type: Option<MemoryType>,
120        limit: usize,
121    ) -> anyhow::Result<Vec<MemoryEntry>> {
122        self.memory_manager.search(query, memory_type, limit).await
123    }
124
125    /// Semantic search using HNSW index.
126    ///
127    /// Falls back to regular search if HNSW index is not available.
128    pub async fn semantic_search_memory(
129        &self,
130        query: &str,
131        memory_type: Option<MemoryType>,
132        limit: usize,
133    ) -> anyhow::Result<Vec<SemanticHit>> {
134        if let Some(ref hnsw) = self.hnsw_index {
135            self.memory_manager
136                .semantic_search(query, memory_type, limit, hnsw)
137                .await
138        } else {
139            // Fallback to regular search, wrap results
140            let entries = self.search_memory(query, memory_type, limit).await?;
141            Ok(entries
142                .into_iter()
143                .map(|entry| SemanticHit {
144                    entry,
145                    distance: 0.0,
146                    similarity: 0.0,
147                })
148                .collect())
149        }
150    }
151
152    /// Memory manager reference.
153    pub fn memory_manager(&self) -> &Arc<MemoryManager> {
154        &self.memory_manager
155    }
156
157    /// Rebuild the HNSW index from all stored memories.
158    pub async fn rebuild_hnsw_index(&self) -> anyhow::Result<usize> {
159        if let Some(ref hnsw) = self.hnsw_index {
160            self.memory_manager.rebuild_hnsw_index(hnsw).await
161        } else {
162            Err(anyhow::anyhow!("HNSW index not initialized"))
163        }
164    }
165}
oxios_kernel/kernel_handle/agent_api.rs

oxios_kernel/kernel_handle/
agent_api.rs
