Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 100 additions & 2 deletions core/src/agent/tools/stimulate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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> {
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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 { .. }));
}
}
105 changes: 105 additions & 0 deletions core/src/engine/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
Loading