Skip to main content

somatize_runtime/
graph_session.rs

1//! Graph session — the primary orchestrator for Graph → Compile → Execute.
2//!
3//! [`GraphSession`] binds a [`Graph`] with its [`NodeCatalog`], cache,
4//! event bus, and optional distributed components into a single object
5//! that can compile, fit, and execute.
6
7use crate::cache::MemoryCache;
8use crate::event_bus::EventBus;
9use crate::executor::{self, Context, GraphInfo};
10use crate::node_catalog::NodeCatalog;
11use crate::runner::Runner;
12use crate::runner::Transport;
13use somatize_compiler::{CompileMode, CompileResult, compile};
14use somatize_core::cache::{CacheKey, CacheStore};
15use somatize_core::error::{Result, SomaError};
16use somatize_core::event::Event;
17use somatize_core::fingerprint::ArchitectureFingerprint;
18use somatize_core::graph::Graph;
19use somatize_core::store::{DataRef, DataStore};
20use somatize_core::util::timestamp_id;
21use somatize_core::value::Value;
22use std::collections::HashMap;
23use std::sync::Arc;
24
25/// The primary orchestrator: Graph + catalog + cache + events.
26///
27/// ```ignore
28/// let mut lib = NodeCatalog::new();
29/// lib.register("scaler", Box::new(MyScaler::new()));
30/// lib.register("model", Box::new(MyModel::new()));
31///
32/// let mut session = GraphSession::new(graph, lib);
33/// session.fit(&train_x, Some(&train_y))?;
34/// let output = session.forward(&test_x)?;
35/// ```
36pub struct GraphSession {
37    graph: Graph,
38    catalog: NodeCatalog,
39    cache: Arc<dyn CacheStore>,
40    event_bus: Arc<EventBus>,
41    data_store: Option<Arc<dyn DataStore>>,
42    transport: Option<Arc<dyn Transport>>,
43    /// Performs and journals step effects. Only needed when the graph
44    /// contains a step; a purely computational graph leaves it unset and
45    /// keeps exactly the old behaviour.
46    driver: Option<crate::effects::EffectDriver>,
47    fitted: bool,
48}
49
50impl GraphSession {
51    /// A session over `graph` with an in-memory cache and its own event
52    /// bus; the `with_*` builders swap in shared or persistent components.
53    pub fn new(graph: Graph, catalog: NodeCatalog) -> Self {
54        Self {
55            graph,
56            catalog,
57            cache: Arc::new(MemoryCache::default()),
58            event_bus: Arc::new(EventBus::new(256)),
59            data_store: None,
60            transport: None,
61            driver: None,
62            fitted: false,
63        }
64    }
65
66    /// Replace the default in-memory cache, e.g. with a tiered or
67    /// persistent store shared across sessions.
68    pub fn with_cache(mut self, cache: Arc<dyn CacheStore>) -> Self {
69        self.cache = cache;
70        self
71    }
72
73    /// Replace the session's own event bus, e.g. with one a tracker
74    /// is already subscribed to.
75    pub fn with_event_bus(mut self, bus: Arc<EventBus>) -> Self {
76        self.event_bus = bus;
77        self
78    }
79
80    /// Attach the data store batched forward passes read rows from.
81    pub fn with_data_store(mut self, store: Arc<dyn DataStore>) -> Self {
82        self.data_store = Some(store);
83        self
84    }
85
86    /// Attach the transport that carries `Remote` plan nodes to workers.
87    pub fn with_transport(mut self, transport: Arc<dyn Transport>) -> Self {
88        self.transport = Some(transport);
89        self
90    }
91
92    /// Attach the effect driver a graph containing steps needs.
93    ///
94    /// The session clones the driver per run and hands it the catalog *at
95    /// that moment*, so filters or steps registered through
96    /// [`Self::catalog_mut`] after this call still count. Without a driver,
97    /// executing a step keeps failing with the executor's own explanation.
98    pub fn with_driver(mut self, driver: crate::effects::EffectDriver) -> Self {
99        self.driver = Some(driver);
100        self
101    }
102
103    /// The stored driver, armed with the catalog as it stands right now.
104    fn run_driver(&self) -> Option<crate::effects::EffectDriver> {
105        self.driver
106            .as_ref()
107            .map(|d| d.clone().with_catalog(Arc::new(self.catalog.clone())))
108    }
109
110    // ── Core operations ──
111
112    /// Compile the graph and return diagnostics without executing.
113    pub fn compile(&self, mode: CompileMode) -> Result<CompileResult> {
114        compile(&self.graph, &self.catalog, mode, Some(self.cache.as_ref()))
115    }
116
117    /// Compile and execute the graph, returning all node outputs.
118    ///
119    /// Emits a `RunStarted`/`RunCompleted` (or `RunFailed`) bracket
120    /// around the node events so readers can compute total duration
121    /// and group the run.
122    pub fn run(&mut self, mode: CompileMode) -> Result<HashMap<String, Value>> {
123        let CompileResult { plan, diagnostics } =
124            compile(&self.graph, &self.catalog, mode, Some(self.cache.as_ref()))?;
125
126        for diag in &diagnostics {
127            tracing::warn!("compile diagnostic: {:?}", diag);
128        }
129
130        let graph_info = GraphInfo::from_graph(&self.graph);
131        let run_id = timestamp_id("graph_run");
132        let mut ctx =
133            Context::new(self.event_bus.clone(), run_id.clone()).with_graph_info(graph_info);
134
135        if let Some(store) = &self.data_store {
136            ctx = ctx.with_data_store(store.clone());
137        }
138        if let Some(transport) = &self.transport {
139            ctx = ctx.with_transport(transport.clone());
140        }
141        if let Some(driver) = self.run_driver() {
142            ctx = ctx.with_driver(driver);
143        }
144
145        self.event_bus.emit(Event::RunStarted {
146            run_id: run_id.clone(),
147            plan_summary: plan.summary(),
148        });
149        let start = std::time::Instant::now();
150        if let Err(e) = executor::execute(&plan, &mut ctx, &self.catalog, self.cache.as_ref()) {
151            self.event_bus.emit(Event::RunFailed {
152                run_id,
153                error: e.to_string(),
154            });
155            return Err(e);
156        }
157        self.event_bus.emit(Event::RunCompleted {
158            run_id,
159            duration: start.elapsed(),
160        });
161
162        Ok(ctx.into_outputs())
163    }
164
165    /// Fit all trainable filters in topological order.
166    /// Delegates to LocalRunner — same execution path as remote workers.
167    ///
168    /// Emits a `RunStarted`/`RunCompleted` (or `RunFailed`) bracket
169    /// tagged with the same run id as the node events inside it.
170    pub fn fit(&mut self, x: &Value, y: Option<&Value>) -> Result<HashMap<String, Value>> {
171        self.graph.validate()?;
172
173        let CompileResult { plan, .. } = compile(
174            &self.graph,
175            &self.catalog,
176            CompileMode::NoCache,
177            Some(self.cache.as_ref()),
178        )?;
179
180        let run_id = timestamp_id("fit");
181        self.event_bus.emit(Event::RunStarted {
182            run_id: run_id.clone(),
183            plan_summary: plan.summary(),
184        });
185        let start = std::time::Instant::now();
186
187        let runner = crate::runner::LocalRunner;
188        let mut ctx = crate::runner::RunContext::new(
189            &self.catalog,
190            self.cache.as_ref(),
191            &self.event_bus,
192            &run_id,
193            GraphInfo::from_graph(&self.graph),
194        );
195        if let Some(driver) = self.run_driver() {
196            ctx = ctx.with_driver(driver);
197        }
198        let result = runner.fit(&plan, &ctx, x, y);
199        let (_last_output, mut all_outputs) = match result {
200            Ok(out) => {
201                self.event_bus.emit(Event::RunCompleted {
202                    run_id,
203                    duration: start.elapsed(),
204                });
205                out
206            }
207            Err(e) => {
208                self.event_bus.emit(Event::RunFailed {
209                    run_id,
210                    error: e.to_string(),
211                });
212                return Err(e);
213            }
214        };
215
216        // Store trained states from __state_ keys into NodeCatalog
217        for (key, value) in &all_outputs {
218            if let Some(node_id) = somatize_core::keys::node_of_state_key(key) {
219                self.catalog.try_set_state(node_id, value.clone())?;
220            }
221        }
222
223        // Remove __state_ keys from returned outputs (callers expect node IDs only)
224        all_outputs.retain(|k, _| somatize_core::keys::node_of_state_key(k).is_none());
225
226        self.fitted = true;
227        Ok(all_outputs)
228    }
229
230    /// Forward pass using the given strategy.
231    ///
232    /// Strategies define HOW data flows through the compiled graph:
233    /// - [`crate::forward::Standard`] — full input at once with inference caching (default)
234    /// - [`crate::forward::Stream`] — chunked input through StreamExecutor
235    /// - [`crate::forward::Batched`] — rows from DataStore, batch by batch
236    pub fn forward_with(
237        &self,
238        x: &Value,
239        strategy: &dyn crate::forward::ForwardStrategy,
240    ) -> Result<Value> {
241        let driver = self.run_driver();
242        strategy.forward(
243            &self.graph,
244            &crate::forward::ForwardEnv {
245                catalog: &self.catalog,
246                cache: self.cache.as_ref(),
247                event_bus: &self.event_bus,
248                data_store: self.data_store.as_ref(),
249                driver: driver.as_ref(),
250            },
251            x,
252        )
253    }
254
255    /// Standard forward pass (shortcut for `forward_with(x, &Standard)`).
256    pub fn forward(&self, x: &Value) -> Result<Value> {
257        self.forward_with(x, &crate::forward::Standard)
258    }
259
260    // ── State persistence ──
261
262    /// Persist all trained states to the data store.
263    pub fn persist_states(&self) -> Result<DataRef> {
264        let store = self
265            .data_store
266            .as_ref()
267            .ok_or_else(|| SomaError::Execution {
268                node_id: "session".into(),
269                message: "persist_states requires a data store".into(),
270            })?;
271
272        let sorted = self.graph.topological_sort()?;
273        let mut states_map = serde_json::Map::new();
274        for node_id in &sorted {
275            if let Some(state) = self.catalog.get_state(node_id) {
276                let json = serde_json::to_value(&*state)
277                    .map_err(|e| SomaError::Other(format!("state serialize: {e}")))?;
278                states_map.insert(node_id.to_string(), json);
279            }
280        }
281
282        let states_value = Value::json(serde_json::Value::Object(states_map));
283        let fingerprint = self.graph_config_hash()?;
284        let key = CacheKey::from_parts(&[b"graph_states", fingerprint.as_bytes()]);
285        store.put(&key, &states_value)
286    }
287
288    /// Load previously persisted states from a data store reference.
289    pub fn load_states(&mut self, data_ref: &DataRef) -> Result<()> {
290        let store = self
291            .data_store
292            .as_ref()
293            .ok_or_else(|| SomaError::Execution {
294                node_id: "session".into(),
295                message: "load_states requires a data store".into(),
296            })?;
297
298        let states_value = store.get(data_ref)?;
299        let states_json = states_value
300            .as_json()
301            .ok_or_else(|| SomaError::Other("persisted states must be JSON".into()))?;
302        let obj = states_json
303            .as_object()
304            .ok_or_else(|| SomaError::Other("persisted states must be a JSON object".into()))?;
305
306        for (node_id, json_val) in obj {
307            let value: Value = serde_json::from_value(json_val.clone())
308                .map_err(|e| SomaError::Other(format!("state deserialize: {e}")))?;
309            self.catalog.try_set_state(node_id.clone(), value)?;
310        }
311
312        self.fitted = true;
313        Ok(())
314    }
315
316    // ── Observability ──
317
318    /// Subscribe to execution events.
319    pub fn subscribe(&self) -> tokio::sync::broadcast::Receiver<Event> {
320        self.event_bus.subscribe()
321    }
322
323    /// Access the event bus directly.
324    pub fn event_bus(&self) -> &Arc<EventBus> {
325        &self.event_bus
326    }
327
328    /// Whether the session has been fitted.
329    pub fn is_fitted(&self) -> bool {
330        self.fitted
331    }
332
333    /// Access the graph.
334    pub fn graph(&self) -> &Graph {
335        &self.graph
336    }
337
338    /// Access the node catalog.
339    pub fn catalog(&self) -> &NodeCatalog {
340        &self.catalog
341    }
342
343    /// Mutable access to the node catalog (for registering nodes after creation).
344    pub fn catalog_mut(&mut self) -> &mut NodeCatalog {
345        &mut self.catalog
346    }
347
348    // ── Private helpers ──
349
350    /// The address under which this graph's trained states are persisted.
351    ///
352    /// It has to follow the graph's *shape*, not just its node names. The
353    /// previous form was `node_ids.join(",")`, so two graphs that shared
354    /// node ids but wired them differently — or configured them
355    /// differently — persisted to one address and read back each other's
356    /// states. [`ArchitectureFingerprint`] already computes exactly this,
357    /// canonically, for the experiment pool.
358    fn graph_config_hash(&self) -> Result<String> {
359        Ok(ArchitectureFingerprint::of(&self.graph)?.digest)
360    }
361}
362
363// ── Convenience free functions ──
364//
365// One-liners over [`GraphSession`]. They used to be separate
366// implementations, and `graph_fit` was the worst of them: a topological
367// loop written from scratch that never compiled a plan, so it ignored
368// parallelism, loops, branches and steps outright — and then discarded
369// every state it fitted instead of storing it. A graph that ran fine
370// through `GraphSession::fit` did something else here.
371
372/// Compile and execute a graph, returning all node outputs.
373pub fn graph_run(
374    graph: &Graph,
375    catalog: &NodeCatalog,
376    mode: CompileMode,
377    cache: Arc<dyn CacheStore>,
378) -> Result<HashMap<String, Value>> {
379    GraphSession::new(graph.clone(), catalog.clone())
380        .with_cache(cache)
381        .run(mode)
382}
383
384/// Fit all trainable filters, returning every node's output.
385pub fn graph_fit(
386    graph: &Graph,
387    catalog: &NodeCatalog,
388    x: &Value,
389    y: Option<&Value>,
390    cache: Arc<dyn CacheStore>,
391) -> Result<HashMap<String, Value>> {
392    GraphSession::new(graph.clone(), catalog.clone())
393        .with_cache(cache)
394        .fit(x, y)
395}
396
397/// Compile in Inference mode and execute, returning the output.
398pub fn graph_predict(
399    graph: &Graph,
400    catalog: &NodeCatalog,
401    x: &Value,
402    cache: Arc<dyn CacheStore>,
403) -> Result<Value> {
404    GraphSession::new(graph.clone(), catalog.clone())
405        .with_cache(cache)
406        .forward(x)
407}
408
409#[cfg(test)]
410mod tests {
411    use super::*;
412    use crate::cache::MemoryCache;
413    use somatize_compiler::NodeRegistry;
414    use somatize_core::cache::CacheKey;
415    use somatize_core::error::Result;
416    use somatize_core::filter::{FilterKind, FilterMeta, StreamMode};
417    use somatize_core::graph::{Edge, Node};
418
419    // ── Test filters ──
420
421    struct DoublerFilter;
422    impl somatize_core::filter::Filter for DoublerFilter {
423        fn config_hash(&self) -> CacheKey {
424            CacheKey::from_parts(&[b"Doubler"])
425        }
426        fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
427            Ok(Value::Empty)
428        }
429        fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
430            let (data, shape) = x
431                .as_tensor()
432                .ok_or(SomaError::Other("need tensor".into()))?;
433            Ok(Value::tensor(
434                data.iter().map(|v| v * 2.0).collect(),
435                shape.to_vec(),
436            ))
437        }
438        fn meta(&self) -> FilterMeta {
439            FilterMeta {
440                name: "Doubler".into(),
441                kind: FilterKind::Stateless,
442                cacheable: true,
443                differentiable: true,
444                deterministic: true,
445                stream_mode: StreamMode::FixedState,
446                distribution: somatize_core::filter::Distribution::Local,
447                input_schema: None,
448                output_schema: None,
449            }
450        }
451    }
452
453    struct AdderFilter(f64);
454    impl somatize_core::filter::Filter for AdderFilter {
455        fn config_hash(&self) -> CacheKey {
456            CacheKey::from_parts(&[b"Adder", &self.0.to_le_bytes()])
457        }
458        fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
459            Ok(Value::Empty)
460        }
461        fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
462            let (data, shape) = x
463                .as_tensor()
464                .ok_or(SomaError::Other("need tensor".into()))?;
465            Ok(Value::tensor(
466                data.iter().map(|v| v + self.0).collect(),
467                shape.to_vec(),
468            ))
469        }
470        fn meta(&self) -> FilterMeta {
471            FilterMeta {
472                name: "Adder".into(),
473                kind: FilterKind::Stateless,
474                cacheable: true,
475                differentiable: true,
476                deterministic: true,
477                stream_mode: StreamMode::FixedState,
478                distribution: somatize_core::filter::Distribution::Local,
479                input_schema: None,
480                output_schema: None,
481            }
482        }
483    }
484
485    struct MeanFilter;
486    impl somatize_core::filter::Filter for MeanFilter {
487        fn config_hash(&self) -> CacheKey {
488            CacheKey::from_parts(&[b"Mean"])
489        }
490        fn fit(&self, x: &Value, _y: Option<&Value>) -> Result<Value> {
491            let (data, _) = x
492                .as_tensor()
493                .ok_or(SomaError::Other("need tensor".into()))?;
494            let mean = data.iter().sum::<f64>() / data.len() as f64;
495            Ok(Value::json(serde_json::json!({ "mean": mean })))
496        }
497        fn forward(&self, x: &Value, state: &Value) -> Result<Value> {
498            let (data, shape) = x
499                .as_tensor()
500                .ok_or(SomaError::Other("need tensor".into()))?;
501            let mean = state
502                .as_json()
503                .and_then(|j| j["mean"].as_f64())
504                .unwrap_or(0.0);
505            Ok(Value::tensor(
506                data.iter().map(|v| v - mean).collect(),
507                shape.to_vec(),
508            ))
509        }
510        fn meta(&self) -> FilterMeta {
511            FilterMeta {
512                name: "Mean".into(),
513                kind: FilterKind::Trainable,
514                cacheable: true,
515                differentiable: true,
516                deterministic: true,
517                stream_mode: StreamMode::FixedState,
518                distribution: somatize_core::filter::Distribution::Local,
519                input_schema: None,
520                output_schema: None,
521            }
522        }
523    }
524
525    fn linear_graph(ids: &[&str]) -> Graph {
526        let mut g = Graph::new();
527        for &id in ids {
528            g.nodes.push(Node::new(id, id, id));
529        }
530        for (i, pair) in ids.windows(2).enumerate() {
531            g.edges.push(Edge::data(format!("e{i}"), pair[0], pair[1]));
532        }
533        g
534    }
535
536    // ── GraphSession tests ──
537
538    #[test]
539    fn session_run_linear() {
540        let graph = linear_graph(&["double", "add"]);
541        let mut lib = NodeCatalog::new();
542        lib.register("double", Box::new(DoublerFilter));
543        lib.register("add", Box::new(AdderFilter(10.0)));
544
545        let mut session = GraphSession::new(graph, lib);
546
547        let cache = MemoryCache::default();
548        session = session.with_cache(Arc::new(cache));
549
550        // Manual compile + execute via run
551        let CompileResult { plan, .. } = session.compile(CompileMode::NoCache).unwrap();
552        let bus = Arc::new(EventBus::new(64));
553        let mut ctx =
554            Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(session.graph()));
555        ctx.set(
556            somatize_core::keys::GRAPH_INPUT,
557            Value::tensor(vec![1.0, 2.0, 3.0], vec![3]),
558        );
559        executor::execute(&plan, &mut ctx, session.catalog(), &MemoryCache::default()).unwrap();
560
561        let outputs: HashMap<String, Value> = ctx.into_outputs();
562
563        let result = outputs.get("add").unwrap();
564        let (data, _) = result.as_tensor().unwrap();
565        assert_eq!(data, &[12.0, 14.0, 16.0]);
566    }
567
568    #[test]
569    fn session_fit_and_forward() {
570        let graph = linear_graph(&["mean", "double"]);
571        let mut lib = NodeCatalog::new();
572        lib.register("mean", Box::new(MeanFilter));
573        lib.register("double", Box::new(DoublerFilter));
574
575        let mut session = GraphSession::new(graph, lib);
576
577        let x = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
578        let outputs = session.fit(&x, None).unwrap();
579
580        // mean: fit learns mean=20, forward: [10-20, 20-20, 30-20] = [-10, 0, 10]
581        // double: [-10, 0, 10] → [-20, 0, 20]
582        let result = outputs.get("double").unwrap();
583        let (data, _) = result.as_tensor().unwrap();
584        assert_eq!(data, &[-20.0, 0.0, 20.0]);
585
586        assert!(session.is_fitted());
587    }
588
589    #[test]
590    fn session_compile_diagnostics() {
591        let graph = linear_graph(&["double"]);
592        let mut lib = NodeCatalog::new();
593        lib.register("double", Box::new(DoublerFilter));
594
595        let session = GraphSession::new(graph, lib);
596        let result = session.compile(CompileMode::NoCache).unwrap();
597        assert!(result.plan.node_count() > 0);
598    }
599
600    // ── Free function tests (backward compat) ──
601
602    #[test]
603    fn graph_run_linear() {
604        let graph = linear_graph(&["double", "add"]);
605        let mut lib = NodeCatalog::new();
606        lib.register("double", Box::new(DoublerFilter));
607        lib.register("add", Box::new(AdderFilter(10.0)));
608
609        let cache = MemoryCache::default();
610
611        let outputs = {
612            let CompileResult { plan, .. } =
613                compile(&graph, &lib, CompileMode::NoCache, None).unwrap();
614            let bus = Arc::new(EventBus::new(64));
615            let mut ctx = Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(&graph));
616            ctx.set(
617                somatize_core::keys::GRAPH_INPUT,
618                Value::tensor(vec![1.0, 2.0, 3.0], vec![3]),
619            );
620            executor::execute(&plan, &mut ctx, &lib, &cache).unwrap();
621            ctx.into_outputs()
622        };
623
624        let result = outputs.get("add").unwrap();
625        let (data, _) = result.as_tensor().unwrap();
626        assert_eq!(data, &[12.0, 14.0, 16.0]);
627    }
628
629    #[test]
630    fn graph_run_diamond() {
631        let mut graph = Graph::new();
632        graph.nodes.push(Node::new("double", "Double", "double"));
633        graph.nodes.push(Node::new("add", "Add", "add"));
634        graph.nodes.push(Node::new("merge", "Merge", "merge"));
635        graph.edges.push(Edge::data("e1", "double", "merge"));
636        graph.edges.push(Edge::data("e2", "add", "merge"));
637
638        let mut lib = NodeCatalog::new();
639        lib.register("double", Box::new(DoublerFilter));
640        lib.register("add", Box::new(AdderFilter(100.0)));
641
642        struct MergeFilter;
643        impl somatize_core::filter::Filter for MergeFilter {
644            fn config_hash(&self) -> CacheKey {
645                CacheKey::from_parts(&[b"Merge"])
646            }
647            fn fit(&self, _x: &Value, _y: Option<&Value>) -> Result<Value> {
648                Ok(Value::Empty)
649            }
650            fn forward(&self, x: &Value, _state: &Value) -> Result<Value> {
651                Ok(x.clone())
652            }
653            fn meta(&self) -> FilterMeta {
654                FilterMeta {
655                    name: "Merge".into(),
656                    kind: FilterKind::Stateless,
657                    cacheable: true,
658                    differentiable: false,
659                    deterministic: true,
660                    stream_mode: StreamMode::FixedState,
661                    distribution: somatize_core::filter::Distribution::Local,
662                    input_schema: None,
663                    output_schema: None,
664                }
665            }
666        }
667        lib.register("merge", Box::new(MergeFilter));
668
669        let cache = MemoryCache::default();
670        let CompileResult { plan, .. } = compile(&graph, &lib, CompileMode::NoCache, None).unwrap();
671
672        let bus = Arc::new(EventBus::new(64));
673        let mut ctx = Context::new(bus, "test").with_graph_info(GraphInfo::from_graph(&graph));
674        ctx.set(
675            somatize_core::keys::GRAPH_INPUT,
676            Value::tensor(vec![5.0], vec![1]),
677        );
678        executor::execute(&plan, &mut ctx, &lib, &cache).unwrap();
679
680        let merge_output = ctx.get("merge").unwrap();
681        assert!(
682            merge_output.as_json().is_some(),
683            "merge should receive JSON from multiple predecessors"
684        );
685    }
686
687    #[test]
688    fn graph_fit_trainable() {
689        let graph = linear_graph(&["mean", "double"]);
690        let mut lib = NodeCatalog::new();
691        lib.register("mean", Box::new(MeanFilter));
692        lib.register("double", Box::new(DoublerFilter));
693
694        let cache = Arc::new(MemoryCache::default());
695        let x = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
696
697        let outputs = graph_fit(&graph, &lib, &x, None, cache.clone()).unwrap();
698
699        let result = outputs.get("double").unwrap();
700        let (data, _) = result.as_tensor().unwrap();
701        assert_eq!(data, &[-20.0, 0.0, 20.0]);
702
703        assert!(!cache.is_empty());
704    }
705
706    #[test]
707    fn the_catalog_is_the_compiler_registry() {
708        let mut lib = NodeCatalog::new();
709        lib.register("a", Box::new(DoublerFilter));
710
711        let registry: &dyn NodeRegistry = &lib;
712        assert!(registry.meta("a").is_some());
713        assert_eq!(registry.meta("a").unwrap().name, "Doubler");
714        assert!(registry.config_hash("a").is_some());
715        assert!(registry.meta("b").is_none());
716    }
717
718    fn session_of(graph: Graph) -> GraphSession {
719        let mut lib = NodeCatalog::new();
720        for node in &graph.nodes {
721            lib.register(&node.id, Box::new(DoublerFilter));
722        }
723        GraphSession::new(graph, lib)
724    }
725
726    /// The persisted-state address follows the wiring, not just the names.
727    /// It used to be `node_ids.join(",")`, so these two graphs shared one
728    /// address and each would load back the other's trained states.
729    #[test]
730    fn state_address_separates_graphs_that_share_node_ids() {
731        let chain = session_of(linear_graph(&["a", "b", "c"]));
732
733        // Same three nodes, different wiring: a fan-out from `a`.
734        let mut fan = Graph::new();
735        for id in ["a", "b", "c"] {
736            fan.nodes.push(Node::new(id, id, id));
737        }
738        fan.edges.push(Edge::data("e0", "a", "b"));
739        fan.edges.push(Edge::data("e1", "a", "c"));
740        let fan = session_of(fan);
741
742        assert_ne!(
743            chain.graph_config_hash().unwrap(),
744            fan.graph_config_hash().unwrap(),
745            "two differently wired graphs must not persist states to one address"
746        );
747    }
748
749    /// The same graph built twice is the same address — otherwise nothing
750    /// persisted could ever be loaded back.
751    #[test]
752    fn state_address_is_stable_for_the_same_graph() {
753        assert_eq!(
754            session_of(linear_graph(&["a", "b"]))
755                .graph_config_hash()
756                .unwrap(),
757            session_of(linear_graph(&["a", "b"]))
758                .graph_config_hash()
759                .unwrap()
760        );
761    }
762}