diff --git a/crates/zclaw-kernel/src/kernel.rs b/crates/zclaw-kernel/src/kernel.rs index df871da..bcc21f1 100644 --- a/crates/zclaw-kernel/src/kernel.rs +++ b/crates/zclaw-kernel/src/kernel.rs @@ -534,6 +534,9 @@ impl Kernel { // Run the loop let result = loop_runner.run(session_id, message).await?; + // Track message count + self.registry.increment_message_count(agent_id); + Ok(MessageResponse { content: result.response, input_tokens: result.input_tokens, @@ -624,6 +627,7 @@ impl Kernel { let loop_runner = loop_runner.with_system_prompt(&system_prompt); // Run with streaming + self.registry.increment_message_count(agent_id); loop_runner.run_streaming(session_id, message).await } diff --git a/crates/zclaw-kernel/src/registry.rs b/crates/zclaw-kernel/src/registry.rs index f78394f..c157071 100644 --- a/crates/zclaw-kernel/src/registry.rs +++ b/crates/zclaw-kernel/src/registry.rs @@ -9,6 +9,7 @@ pub struct AgentRegistry { agents: DashMap, states: DashMap, created_at: DashMap>, + message_counts: DashMap, } impl AgentRegistry { @@ -17,6 +18,7 @@ impl AgentRegistry { agents: DashMap::new(), states: DashMap::new(), created_at: DashMap::new(), + message_counts: DashMap::new(), } } @@ -33,6 +35,7 @@ impl AgentRegistry { self.agents.remove(id); self.states.remove(id); self.created_at.remove(id); + self.message_counts.remove(id); } /// Get an agent by ID @@ -53,7 +56,7 @@ impl AgentRegistry { model: config.model.model.clone(), provider: config.model.provider.clone(), state, - message_count: 0, // TODO: Track this + message_count: self.message_counts.get(id).map(|c| *c as usize).unwrap_or(0), created_at, updated_at: Utc::now(), }) @@ -83,6 +86,11 @@ impl AgentRegistry { pub fn count(&self) -> usize { self.agents.len() } + + /// Increment message count for an agent + pub fn increment_message_count(&self, id: &AgentId) { + self.message_counts.entry(*id).and_modify(|c| *c += 1).or_insert(1); + } } impl Default for AgentRegistry {