From d593a2198bd5d2a9b7be2c8835a57b9bab421952 Mon Sep 17 00:00:00 2001 From: ammar siddiqui Date: Wed, 27 May 2026 17:09:00 -0400 Subject: [PATCH] Add stimulation agent tools --- core/src/agent/registry.rs | 2 + core/src/agent/tools/mod.rs | 1 + core/src/agent/tools/stimulate.rs | 227 ++++++++++++++++++++++++++++++ core/src/engine/mod.rs | 30 ++++ 4 files changed, 260 insertions(+) create mode 100644 core/src/agent/tools/stimulate.rs diff --git a/core/src/agent/registry.rs b/core/src/agent/registry.rs index c7fd7f1..c4827cb 100644 --- a/core/src/agent/registry.rs +++ b/core/src/agent/registry.rs @@ -23,6 +23,7 @@ use std::collections::BTreeMap; use super::{AgentTool, Permission, ToolContext, ToolDescriptor, ToolError}; use crate::agent::tools::read_only::read_only_tools; +use crate::agent::tools::stimulate::stimulate_tools; use crate::agent::tools::topology_write::topology_write_tools; /// Errors a registry surfaces from operations other than `invoke()`. @@ -101,6 +102,7 @@ impl Registry { pub fn register_builtin_tools(&mut self) -> Result<(), RegistryError> { for tool in read_only_tools() .into_iter() + .chain(stimulate_tools().into_iter()) .chain(topology_write_tools().into_iter()) { self.register(tool)?; diff --git a/core/src/agent/tools/mod.rs b/core/src/agent/tools/mod.rs index 8e3b4f9..25aba23 100644 --- a/core/src/agent/tools/mod.rs +++ b/core/src/agent/tools/mod.rs @@ -1,4 +1,5 @@ //! Built-in agent tools. pub mod read_only; +pub mod stimulate; pub mod topology_write; diff --git a/core/src/agent/tools/stimulate.rs b/core/src/agent/tools/stimulate.rs new file mode 100644 index 0000000..53aecb6 --- /dev/null +++ b/core/src/agent/tools/stimulate.rs @@ -0,0 +1,227 @@ +//! Stimulation tools backed by the live engine. + +use async_trait::async_trait; +use uuid::Uuid; + +use crate::agent::{AgentTool, Permission, ToolContext, ToolDescriptor, ToolError}; + +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; + +pub struct InjectCurrentTool; +pub struct ForceSpikeTool; + +#[async_trait] +impl AgentTool for InjectCurrentTool { + fn descriptor(&self) -> ToolDescriptor { + ToolDescriptor { + name: "inject_current", + description: "Inject clamped current into one neuron for a clamped duration.", + permission: Permission::Stimulate, + input_schema: serde_json::json!({ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "node_id": { "type": "string", "format": "uuid" }, + "current_pa": { "type": "number" }, + "duration_ms": { "type": "number", "default": 20.0 } + }, + "required": ["node_id", "current_pa"], + "additionalProperties": false, + }), + } + } + + async fn invoke( + &self, + args: serde_json::Value, + ctx: ToolContext, + ) -> Result { + let node_id = parse_uuid_arg(&args, "node_id")?; + let requested_current = parse_f32_arg(&args, "current_pa")?; + let requested_duration = args + .get("duration_ms") + .map(|_| parse_f32_arg(&args, "duration_ms")) + .transpose()? + .unwrap_or(20.0); + let current_pa = requested_current.clamp(MIN_CURRENT_PA, MAX_CURRENT_PA); + let duration_ms = requested_duration.clamp(MIN_DURATION_MS, MAX_DURATION_MS); + + ctx.sim()? + .stimulate(node_id, current_pa, duration_ms) + .await + .map_err(|msg| ToolError::Substrate(msg.into()))?; + Ok(serde_json::json!({ + "node_id": node_id, + "requested_current_pa": requested_current, + "current_pa": current_pa, + "requested_duration_ms": requested_duration, + "duration_ms": duration_ms, + })) + } +} + +#[async_trait] +impl AgentTool for ForceSpikeTool { + fn descriptor(&self) -> ToolDescriptor { + ToolDescriptor { + name: "force_spike", + description: "Force one live neuron to emit an immediate spike frame.", + permission: Permission::Stimulate, + input_schema: serde_json::json!({ + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "node_id": { "type": "string", "format": "uuid" } + }, + "required": ["node_id"], + "additionalProperties": false, + }), + } + } + + async fn invoke( + &self, + args: serde_json::Value, + ctx: ToolContext, + ) -> Result { + let node_id = parse_uuid_arg(&args, "node_id")?; + ctx.sim()? + .force_spike(node_id) + .await + .map_err(ToolError::Substrate)?; + Ok(serde_json::json!({ "node_id": node_id, "forced": true })) + } +} + +pub fn stimulate_tools() -> Vec> { + vec![Box::new(InjectCurrentTool), Box::new(ForceSpikeTool)] +} + +fn parse_uuid_arg(args: &serde_json::Value, property: &'static str) -> Result { + let raw = args + .get(property) + .and_then(|value| value.as_str()) + .ok_or_else(|| bad_input(property, format!("{property} must be a UUID string")))?; + Uuid::parse_str(raw).map_err(|e| bad_input(property, format!("{property}: {e}"))) +} + +fn parse_f32_arg(args: &serde_json::Value, property: &'static str) -> Result { + let value = args + .get(property) + .and_then(|value| value.as_f64()) + .ok_or_else(|| bad_input(property, format!("{property} must be a finite number")))?; + let value = value as f32; + if value.is_finite() { + Ok(value) + } else { + Err(bad_input(property, format!("{property} must be finite"))) + } +} + +fn bad_input(tool: &'static str, message: impl Into) -> ToolError { + ToolError::BadInput { + tool: tool.into(), + errors: vec![message.into()], + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::Registry; + use crate::engine::spawn_engine; + use tokio::time::{timeout, Duration}; + + fn registry() -> Registry { + let mut registry = Registry::new(); + for tool in stimulate_tools() { + registry.register(tool).unwrap(); + } + registry + } + + #[tokio::test] + async fn stimulate_tools_register_with_stimulate_permission() { + let registry = registry(); + let descriptors = registry.descriptors(); + assert_eq!( + descriptors.iter().map(|d| d.name).collect::>(), + vec!["force_spike", "inject_current"] + ); + assert!(descriptors + .iter() + .all(|descriptor| descriptor.permission == Permission::Stimulate)); + } + + #[tokio::test] + async fn inject_current_clamps_current_and_duration() { + let (sim, _join) = spawn_engine(10_000); + let node_id = Uuid::new_v4(); + sim.add_node(node_id).await.unwrap(); + + let out = registry() + .invoke( + "inject_current", + serde_json::json!({ + "node_id": node_id, + "current_pa": 250.0, + "duration_ms": 5000.0 + }), + Permission::Stimulate, + ToolContext::new(sim), + ) + .await + .unwrap(); + + assert_eq!(out["requested_current_pa"], 250.0); + assert_eq!(out["current_pa"], 100.0); + assert_eq!(out["requested_duration_ms"], 5000.0); + assert_eq!(out["duration_ms"], 1000.0); + } + + #[tokio::test] + async fn force_spike_emits_observable_spike_frame() { + let (sim, _join) = spawn_engine(10_000); + let node_id = Uuid::new_v4(); + sim.add_node(node_id).await.unwrap(); + let mut spikes = sim.spikes.subscribe(); + + registry() + .invoke( + "force_spike", + serde_json::json!({ "node_id": node_id }), + Permission::Stimulate, + ToolContext::new(sim), + ) + .await + .unwrap(); + + let frame = timeout(Duration::from_millis(200), spikes.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(frame.events.len(), 1); + assert_eq!(frame.events[0].node_id, node_id); + } + + #[tokio::test] + async fn stimulate_tools_require_stimulate_permission() { + let (sim, _join) = spawn_engine(10_000); + let node_id = Uuid::new_v4(); + sim.add_node(node_id).await.unwrap(); + + let err = registry() + .invoke( + "force_spike", + serde_json::json!({ "node_id": node_id }), + Permission::TopologyWrite, + ToolContext::new(sim), + ) + .await + .unwrap_err(); + assert!(matches!(err, ToolError::PermissionDenied { .. })); + } +} diff --git a/core/src/engine/mod.rs b/core/src/engine/mod.rs index 7505267..f75f4dc 100644 --- a/core/src/engine/mod.rs +++ b/core/src/engine/mod.rs @@ -96,6 +96,10 @@ pub enum EngineCommand { current: f32, duration_ms: f32, }, + ForceSpike { + node_id: Uuid, + reply: oneshot::Sender>, + }, Snapshot(oneshot::Sender), /// Live topology snapshot with synapse endpoints. This is the /// read-only graph view the agent harness needs without going @@ -325,6 +329,15 @@ impl SimHandle { .map_err(|_| "engine offline") } + pub async fn force_spike(&self, node_id: Uuid) -> Result<(), String> { + let (tx, rx) = oneshot::channel(); + self.cmd_tx + .send(EngineCommand::ForceSpike { node_id, reply: tx }) + .await + .map_err(|_| "engine offline".to_string())?; + rx.await.map_err(|_| "engine dropped reply".to_string())? + } + pub async fn snapshot(&self) -> Result { let (tx, rx) = oneshot::channel(); self.cmd_tx @@ -707,6 +720,23 @@ pub fn spawn_engine(tick_hz: u32) -> (SimHandle, tokio::task::JoinHandle<()>) { } EngineCommand::Stimulate { node_id, current, duration_ms } => engine.inject(node_id, current, duration_ms), + EngineCommand::ForceSpike { node_id, reply } => { + let result = if engine.neurons.contains_key(&node_id) { + engine.fired_prev.insert(node_id); + let frame = SpikeFrame::new( + engine.t_ms, + vec![hebb::engine::events::SpikeEvent { + node_id, + t_ms: engine.t_ms, + }], + ); + let _ = spike_tx.send(frame); + Ok(()) + } else { + Err(format!("node {node_id} not in engine")) + }; + let _ = reply.send(result); + } EngineCommand::Snapshot(reply) => { let snap = EngineSnapshot { t_ms: engine.t_ms,