Add run_for agent tool (#54)

Author: ammar-s847

core/src/agent/tools/stimulate.rs

@@ -9,9 +9,11 @@ const MIN_CURRENT_PA: f32 = -100.0;
 const MAX_CURRENT_PA: f32 = 100.0;
 const MIN_DURATION_MS: f32 = 0.0;
 const MAX_DURATION_MS: f32 = 1000.0;
+const MAX_RUN_DURATION_MS: f32 = 5000.0;
 
 pub struct InjectCurrentTool;
 pub struct ForceSpikeTool;
+pub struct RunForTool;
 
 #[async_trait]
 impl AgentTool for InjectCurrentTool {
@@ -96,8 +98,64 @@ impl AgentTool for ForceSpikeTool {
     }
 }
 
+#[async_trait]
+impl AgentTool for RunForTool {
+    fn descriptor(&self) -> ToolDescriptor {
+        ToolDescriptor {
+            name: "run_for",
+            description:
+                "Advance the simulator for a bounded duration and return a spike-count summary.",
+            permission: Permission::Stimulate,
+            input_schema: serde_json::json!({
+                "$schema": "http://json-schema.org/draft-07/schema#",
+                "type": "object",
+                "properties": {
+                    "duration_ms": {
+                        "type": "number",
+                        "minimum": 0.0,
+                        "maximum": MAX_RUN_DURATION_MS
+                    },
+                    "dt_ms": {
+                        "type": "number",
+                        "exclusiveMinimum": 0.0,
+                        "maximum": 100.0,
+                        "default": 1.0
+                    }
+                },
+                "required": ["duration_ms"],
+                "additionalProperties": false,
+            }),
+        }
+    }
+
+    async fn invoke(
+        &self,
+        args: serde_json::Value,
+        ctx: ToolContext,
+    ) -> Result<serde_json::Value, ToolError> {
+        let requested_duration = parse_f32_arg(&args, "duration_ms")?;
+        let requested_dt = args
+            .get("dt_ms")
+            .map(|_| parse_f32_arg(&args, "dt_ms"))
+            .transpose()?
+            .unwrap_or(1.0);
+        let duration_ms = requested_duration.clamp(0.0, MAX_RUN_DURATION_MS);
+        let dt_ms = requested_dt.clamp(0.001, 100.0);
+        let summary = ctx
+            .sim()?
+            .run_for(duration_ms, dt_ms)
+            .await
+            .map_err(ToolError::Substrate)?;
+        serde_json::to_value(summary).map_err(|e| ToolError::Substrate(e.to_string()))
+    }
+}
+
 pub fn stimulate_tools() -> Vec<Box<dyn AgentTool>> {
-    vec![Box::new(InjectCurrentTool), Box::new(ForceSpikeTool)]
+    vec![
+        Box::new(InjectCurrentTool),
+        Box::new(ForceSpikeTool),
+        Box::new(RunForTool),
+    ]
 }
 
 fn parse_uuid_arg(args: &serde_json::Value, property: &'static str) -> Result<Uuid, ToolError> {
@@ -149,7 +207,7 @@ mod tests {
         let descriptors = registry.descriptors();
         assert_eq!(
             descriptors.iter().map(|d| d.name).collect::<Vec<_>>(),
-            vec!["force_spike", "inject_current"]
+            vec!["force_spike", "inject_current", "run_for"]
         );
         assert!(descriptors
             .iter()
@@ -224,4 +282,44 @@ mod tests {
             .unwrap_err();
         assert!(matches!(err, ToolError::PermissionDenied { .. }));
     }
+
+    #[tokio::test]
+    async fn run_for_returns_structured_spike_summary() {
+        let (sim, _join) = spawn_engine(10_000);
+        let node_id = Uuid::new_v4();
+        sim.add_node(node_id).await.unwrap();
+        sim.force_spike(node_id).await.unwrap();
+
+        let out = registry()
+            .invoke(
+                "run_for",
+                serde_json::json!({ "duration_ms": 500.0, "dt_ms": 1.0 }),
+                Permission::Stimulate,
+                ToolContext::new(sim),
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(out["duration_ms"], 500.0);
+        assert_eq!(out["dt_ms"], 1.0);
+        assert_eq!(out["steps"], 500);
+        assert!(out["total_spikes"].as_u64().is_some());
+        assert_eq!(out["cancelled"], false);
+        assert!(out["per_neuron"].as_array().is_some());
+    }
+
+    #[tokio::test]
+    async fn run_for_is_bounded() {
+        let (sim, _join) = spawn_engine(10_000);
+        let out = registry()
+            .invoke(
+                "run_for",
+                serde_json::json!({ "duration_ms": 6000.0, "dt_ms": 10.0 }),
+                Permission::Stimulate,
+                ToolContext::new(sim),
+            )
+            .await
+            .unwrap_err();
+        assert!(matches!(out, ToolError::BadInput { .. }));
+    }
 }

core/src/engine/mod.rs

@@ -20,6 +20,7 @@ pub use hebb::engine::sim::SimEngine;
 pub use spike_persist::spawn_spike_persister;
 pub use weight_persist::spawn_weight_persister;
 
+use std::collections::BTreeMap;
 use std::path::PathBuf;
 use std::time::Duration;
 use tokio::sync::{broadcast, mpsc, oneshot};
@@ -73,6 +74,25 @@ pub struct EngineSeedReport {
     pub added_edges: usize,
 }
 
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct EngineRunSummary {
+    pub requested_duration_ms: f32,
+    pub duration_ms: f32,
+    pub dt_ms: f32,
+    pub steps: usize,
+    pub t_start_ms: f64,
+    pub t_end_ms: f64,
+    pub total_spikes: usize,
+    pub per_neuron: Vec<EngineRunNeuronCount>,
+    pub cancelled: bool,
+}
+
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct EngineRunNeuronCount {
+    pub node_id: Uuid,
+    pub spikes: usize,
+}
+
 pub enum EngineCommand {
     AddNode(Uuid),
     AddEdge {
@@ -100,6 +120,11 @@ pub enum EngineCommand {
         node_id: Uuid,
         reply: oneshot::Sender<Result<(), String>>,
     },
+    RunFor {
+        duration_ms: f32,
+        dt_ms: f32,
+        reply: oneshot::Sender<Result<EngineRunSummary, String>>,
+    },
     Snapshot(oneshot::Sender<EngineSnapshot>),
     /// Live topology snapshot with synapse endpoints. This is the
     /// read-only graph view the agent harness needs without going
@@ -338,6 +363,19 @@ impl SimHandle {
         rx.await.map_err(|_| "engine dropped reply".to_string())?
     }
 
+    pub async fn run_for(&self, duration_ms: f32, dt_ms: f32) -> Result<EngineRunSummary, String> {
+        let (tx, rx) = oneshot::channel();
+        self.cmd_tx
+            .send(EngineCommand::RunFor {
+                duration_ms,
+                dt_ms,
+                reply: tx,
+            })
+            .await
+            .map_err(|_| "engine offline".to_string())?;
+        rx.await.map_err(|_| "engine dropped reply".to_string())?
+    }
+
     pub async fn snapshot(&self) -> Result<EngineSnapshot, &'static str> {
         let (tx, rx) = oneshot::channel();
         self.cmd_tx
@@ -737,6 +775,19 @@ pub fn spawn_engine(tick_hz: u32) -> (SimHandle, tokio::task::JoinHandle<()>) {
                             };
                             let _ = reply.send(result);
                         }
+                        EngineCommand::RunFor {
+                            duration_ms,
+                            dt_ms,
+                            reply,
+                        } => {
+                            let result = run_engine_for(
+                                &mut engine,
+                                &spike_tx,
+                                duration_ms,
+                                dt_ms,
+                            );
+                            let _ = reply.send(result);
+                        }
                         EngineCommand::Snapshot(reply) => {
                             let snap = EngineSnapshot {
                                 t_ms: engine.t_ms,
@@ -1830,6 +1881,60 @@ fn open_folder_into_engine(
     Ok((engine, cortex_type, kind, cortex, summary))
 }
 
+fn run_engine_for(
+    engine: &mut SimEngine,
+    spike_tx: &broadcast::Sender<SpikeFrame>,
+    duration_ms: f32,
+    dt_ms: f32,
+) -> Result<EngineRunSummary, String> {
+    if !duration_ms.is_finite() || !dt_ms.is_finite() {
+        return Err("duration_ms and dt_ms must be finite".into());
+    }
+    if duration_ms < 0.0 {
+        return Err("duration_ms must be >= 0".into());
+    }
+    if dt_ms <= 0.0 {
+        return Err("dt_ms must be > 0".into());
+    }
+
+    let t_start_ms = engine.t_ms;
+    let steps = (duration_ms / dt_ms).ceil() as usize;
+    let mut counts: BTreeMap<Uuid, usize> = BTreeMap::new();
+    let mut total_spikes = 0;
+
+    for step in 0..steps {
+        let elapsed = step as f32 * dt_ms;
+        let remaining = (duration_ms - elapsed).max(0.0);
+        let step_dt = remaining.min(dt_ms);
+        if step_dt <= 0.0 {
+            break;
+        }
+        let frame = engine.tick(step_dt);
+        if !frame.events.is_empty() {
+            total_spikes += frame.events.len();
+            for event in &frame.events {
+                *counts.entry(event.node_id).or_insert(0) += 1;
+            }
+            let _ = spike_tx.send(frame);
+        }
+    }
+
+    Ok(EngineRunSummary {
+        requested_duration_ms: duration_ms,
+        duration_ms,
+        dt_ms,
+        steps,
+        t_start_ms,
+        t_end_ms: engine.t_ms,
+        total_spikes,
+        per_neuron: counts
+            .into_iter()
+            .map(|(node_id, spikes)| EngineRunNeuronCount { node_id, spikes })
+            .collect(),
+        cancelled: false,
+    })
+}
+
 /// Periodically diff the live weight set against the last published one
 /// and broadcast deltas. Runs as a sibling task to the engine — reads
 /// state through `SimHandle::weight_snapshot` (one mpsc round-trip) so