dynamo_runtime/pipeline/
context.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::BTreeMap;
5use std::ops::{Deref, DerefMut};
6use std::sync::{Arc, Mutex};
7
8use super::{AsyncEngineContext, AsyncEngineContextProvider, Data};
9use crate::engine::{AsyncEngineController, EngineContextGuard};
10use async_trait::async_trait;
11
12use super::registry::Registry;
13
14pub struct Context<T: Data> {
15    current: T,
16    controller: Arc<Controller>, //todo: hold this as an arc
17    registry: Registry,
18    stages: Vec<String>,
19    metadata: BTreeMap<String, String>,
20}
21
22impl<T: Send + Sync + 'static> Context<T> {
23    // Create a new context with initial data
24    pub fn new(current: T) -> Self {
25        Context {
26            current,
27            controller: Arc::new(Controller::default()),
28            registry: Registry::new(),
29            stages: Vec::new(),
30            metadata: BTreeMap::new(),
31        }
32    }
33
34    pub fn rejoin<U: Send + Sync + 'static>(current: T, context: Context<U>) -> Self {
35        Context {
36            current,
37            controller: context.controller,
38            registry: context.registry,
39            stages: context.stages,
40            metadata: context.metadata,
41        }
42    }
43
44    pub fn with_controller(current: T, controller: Controller) -> Self {
45        Context {
46            current,
47            controller: Arc::new(controller),
48            registry: Registry::new(),
49            stages: Vec::new(),
50            metadata: BTreeMap::new(),
51        }
52    }
53
54    pub fn with_id_and_metadata(
55        current: T,
56        id: String,
57        metadata: BTreeMap<String, String>,
58    ) -> Self {
59        Context {
60            current,
61            controller: Arc::new(Controller::new(id)),
62            registry: Registry::new(),
63            stages: Vec::new(),
64            metadata,
65        }
66    }
67
68    /// Get the id of the context
69    pub fn id(&self) -> &str {
70        self.controller.id()
71    }
72
73    /// Get the content of the context
74    pub fn content(&self) -> &T {
75        &self.current
76    }
77
78    pub fn controller(&self) -> &Controller {
79        &self.controller
80    }
81
82    pub fn metadata(&self) -> &BTreeMap<String, String> {
83        &self.metadata
84    }
85
86    pub fn metadata_mut(&mut self) -> &mut BTreeMap<String, String> {
87        &mut self.metadata
88    }
89
90    pub fn set_metadata(&mut self, metadata: BTreeMap<String, String>) {
91        self.metadata = metadata;
92    }
93
94    pub fn insert_metadata<K: Into<String>, V: Into<String>>(&mut self, key: K, value: V) {
95        self.metadata.insert(key.into(), value.into());
96    }
97
98    /// Insert an object into the registry with a specific key.
99    pub fn insert<K: ToString, U: Send + Sync + 'static>(&mut self, key: K, value: U) {
100        self.registry.insert_shared(key, value);
101    }
102
103    /// Insert a unique and takable object into the registry with a specific key.
104    pub fn insert_unique<K: ToString, U: Send + Sync + 'static>(&mut self, key: K, value: U) {
105        self.registry.insert_unique(key, value);
106    }
107
108    /// Retrieve an object from the registry by key and type.
109    pub fn get<V: Send + Sync + 'static>(&self, key: &str) -> Result<Arc<V>, String> {
110        self.registry.get_shared(key)
111    }
112
113    /// Retrieve an optional object from the registry by key and type.
114    pub fn get_optional<V: Send + Sync + 'static>(
115        &self,
116        key: &str,
117    ) -> Result<Option<Arc<V>>, String> {
118        self.registry.get_shared_optional(key)
119    }
120
121    /// Clone a unique object from the registry by key and type.
122    pub fn clone_unique<V: Clone + Send + Sync + 'static>(&self, key: &str) -> Result<V, String> {
123        self.registry.clone_unique(key)
124    }
125
126    /// Take a unique object from the registry by key and type.
127    pub fn take_unique<V: Send + Sync + 'static>(&mut self, key: &str) -> Result<V, String> {
128        self.registry.take_unique(key)
129    }
130
131    /// Transfer the Context to a new Object without updating the registry
132    /// This returns a tuple of the previous object and the new Context
133    pub fn transfer<U: Send + Sync + 'static>(self, new_current: U) -> (T, Context<U>) {
134        (
135            self.current,
136            Context {
137                current: new_current,
138                controller: self.controller,
139                registry: self.registry,
140                stages: self.stages,
141                metadata: self.metadata,
142            },
143        )
144    }
145
146    /// Separate out the current object and context
147    pub fn into_parts(self) -> (T, Context<()>) {
148        self.transfer(())
149    }
150
151    pub fn stages(&self) -> &Vec<String> {
152        &self.stages
153    }
154
155    pub fn add_stage(&mut self, stage: &str) {
156        self.stages.push(stage.to_string());
157    }
158
159    /// Transforms the current context to another type using a provided function.
160    pub fn map<U: Send + Sync + 'static, F>(self, f: F) -> Context<U>
161    where
162        F: FnOnce(T) -> U,
163    {
164        // Use the transfer method to move the current value out
165        let (current, temp_context) = self.transfer(());
166
167        // Apply the transformation function to the current value
168        let new_current = f(current);
169
170        // Use transfer again to create the new context with the transformed type
171        temp_context.transfer(new_current).1
172    }
173
174    pub fn try_map<U, F, E>(self, f: F) -> Result<Context<U>, E>
175    where
176        F: FnOnce(T) -> Result<U, E>,
177        U: Send + Sync + 'static,
178    {
179        // Use the transfer method to move the current value out
180        let (current, temp_context) = self.transfer(());
181
182        // Apply the transformation function to the current value
183        let new_current = f(current)?;
184
185        // Use transfer again to create the new context with the transformed type
186        Ok(temp_context.transfer(new_current).1)
187    }
188}
189
190impl<T: Data> std::fmt::Debug for Context<T> {
191    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192        f.debug_struct("Context")
193            .field("id", &self.controller.id())
194            .finish()
195    }
196}
197
198// Implement Deref to allow Context<T> to act like &T
199impl<T: Data> Deref for Context<T> {
200    type Target = T;
201
202    fn deref(&self) -> &Self::Target {
203        &self.current
204    }
205}
206
207// Implement DerefMut to allow Context<T> to act like &mut T
208impl<T: Data> DerefMut for Context<T> {
209    fn deref_mut(&mut self) -> &mut Self::Target {
210        &mut self.current
211    }
212}
213
214// Implement the custom trait for Context<T>
215impl<T> From<T> for Context<T>
216where
217    T: Send + Sync + 'static,
218{
219    fn from(current: T) -> Self {
220        Context::new(current)
221    }
222}
223
224// Define a custom trait for conversion from Context<T> to Context<U>
225pub trait IntoContext<U: Data> {
226    fn into_context(self) -> Context<U>;
227}
228
229// Implement the custom trait for converting Context<T> to Context<U>
230impl<T, U> IntoContext<U> for Context<T>
231where
232    T: Send + Sync + 'static + Into<U>,
233    U: Send + Sync + 'static,
234{
235    fn into_context(self) -> Context<U> {
236        self.map(|current| current.into())
237    }
238}
239
240impl<T: Data> AsyncEngineContextProvider for Context<T> {
241    fn context(&self) -> Arc<dyn AsyncEngineContext> {
242        self.controller.clone()
243    }
244}
245
246#[derive(Debug, Clone)]
247pub struct StreamContext {
248    controller: Arc<Controller>,
249    registry: Arc<Registry>,
250    stages: Vec<String>,
251    metadata: BTreeMap<String, String>,
252}
253
254impl StreamContext {
255    fn new(
256        controller: Arc<Controller>,
257        registry: Registry,
258        metadata: BTreeMap<String, String>,
259    ) -> Self {
260        StreamContext {
261            controller,
262            registry: Arc::new(registry),
263            stages: Vec::new(),
264            metadata,
265        }
266    }
267
268    /// Retrieve an object from the registry by key and type.
269    pub fn get<V: Send + Sync + 'static>(&self, key: &str) -> Result<Arc<V>, String> {
270        self.registry.get_shared(key)
271    }
272
273    /// Clone a unique object from the registry by key and type.
274    pub fn clone_unique<V: Clone + Send + Sync + 'static>(&self, key: &str) -> Result<V, String> {
275        self.registry.clone_unique(key)
276    }
277
278    pub fn registry(&self) -> Arc<Registry> {
279        self.registry.clone()
280    }
281
282    pub fn stages(&self) -> &Vec<String> {
283        &self.stages
284    }
285
286    pub fn add_stage(&mut self, stage: &str) {
287        self.stages.push(stage.to_string());
288    }
289
290    pub fn metadata(&self) -> &BTreeMap<String, String> {
291        &self.metadata
292    }
293}
294
295#[async_trait]
296impl AsyncEngineContext for StreamContext {
297    fn id(&self) -> &str {
298        self.controller.id()
299    }
300
301    fn stop(&self) {
302        self.controller.stop();
303    }
304
305    fn kill(&self) {
306        self.controller.kill();
307    }
308
309    fn stop_generating(&self) {
310        self.controller.stop_generating();
311    }
312
313    fn is_stopped(&self) -> bool {
314        self.controller.is_stopped()
315    }
316
317    fn is_killed(&self) -> bool {
318        self.controller.is_killed()
319    }
320
321    async fn stopped(&self) {
322        self.controller.stopped().await
323    }
324
325    async fn killed(&self) {
326        self.controller.killed().await
327    }
328
329    fn link_child(&self, child: Arc<dyn AsyncEngineContext>) {
330        self.controller.link_child(child);
331    }
332
333    fn retain(&self, guard: EngineContextGuard) {
334        self.controller.retain(guard);
335    }
336}
337
338impl AsyncEngineContextProvider for StreamContext {
339    fn context(&self) -> Arc<dyn AsyncEngineContext> {
340        self.controller.clone()
341    }
342}
343
344impl<T: Send + Sync + 'static> From<Context<T>> for StreamContext {
345    fn from(value: Context<T>) -> Self {
346        StreamContext::new(value.controller, value.registry, value.metadata)
347    }
348}
349
350// TODO - refactor here - this came from the dynamo.llm-async-engine crate
351
352use tokio::sync::watch::{Receiver, Sender, channel};
353
354#[derive(Debug, Eq, PartialEq)]
355enum State {
356    Live,
357    Stopped,
358    Killed,
359}
360
361/// A context implementation with cancellation propagation.
362#[derive(Debug)]
363pub struct Controller {
364    id: String,
365    tx: Sender<State>,
366    rx: Receiver<State>,
367    child_context: Mutex<Vec<Arc<dyn AsyncEngineContext>>>,
368    retained: Mutex<Vec<RetainedGuard>>,
369}
370
371struct RetainedGuard {
372    _guard: EngineContextGuard,
373}
374
375impl std::fmt::Debug for RetainedGuard {
376    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
377        f.write_str("RetainedGuard")
378    }
379}
380
381impl Controller {
382    pub fn new(id: String) -> Self {
383        let (tx, rx) = channel(State::Live);
384        Self {
385            id,
386            tx,
387            rx,
388            child_context: Mutex::new(Vec::new()),
389            retained: Mutex::new(Vec::new()),
390        }
391    }
392
393    pub fn id(&self) -> &str {
394        &self.id
395    }
396}
397
398impl Default for Controller {
399    fn default() -> Self {
400        Self::new(uuid::Uuid::new_v4().to_string())
401    }
402}
403
404impl AsyncEngineController for Controller {}
405
406#[async_trait]
407impl AsyncEngineContext for Controller {
408    fn id(&self) -> &str {
409        &self.id
410    }
411
412    fn is_stopped(&self) -> bool {
413        *self.rx.borrow() != State::Live
414    }
415
416    fn is_killed(&self) -> bool {
417        *self.rx.borrow() == State::Killed
418    }
419
420    async fn stopped(&self) {
421        let mut rx = self.rx.clone();
422        loop {
423            if *rx.borrow_and_update() != State::Live || rx.changed().await.is_err() {
424                return;
425            }
426        }
427    }
428
429    async fn killed(&self) {
430        let mut rx = self.rx.clone();
431        loop {
432            if *rx.borrow_and_update() == State::Killed || rx.changed().await.is_err() {
433                return;
434            }
435        }
436    }
437
438    fn stop_generating(&self) {
439        // Clone child Arcs to avoid deadlock if parent is accidentally linked under child
440        let children = self
441            .child_context
442            .lock()
443            .expect("Failed to lock child context")
444            .iter()
445            .cloned()
446            .collect::<Vec<_>>();
447        for child in children {
448            child.stop_generating();
449        }
450
451        let _ = self.tx.send(State::Stopped);
452    }
453
454    fn stop(&self) {
455        // Clone child Arcs to avoid deadlock if parent is accidentally linked under child
456        let children = self
457            .child_context
458            .lock()
459            .expect("Failed to lock child context")
460            .iter()
461            .cloned()
462            .collect::<Vec<_>>();
463        for child in children {
464            child.stop();
465        }
466
467        let _ = self.tx.send(State::Stopped);
468    }
469
470    fn kill(&self) {
471        // Clone child Arcs to avoid deadlock if parent is accidentally linked under child
472        let children = self
473            .child_context
474            .lock()
475            .expect("Failed to lock child context")
476            .iter()
477            .cloned()
478            .collect::<Vec<_>>();
479        for child in children {
480            child.kill();
481        }
482
483        let _ = self.tx.send(State::Killed);
484    }
485
486    fn link_child(&self, child: Arc<dyn AsyncEngineContext>) {
487        self.child_context
488            .lock()
489            .expect("Failed to lock child context")
490            .push(child);
491    }
492
493    fn retain(&self, guard: EngineContextGuard) {
494        self.retained
495            .lock()
496            .expect("Failed to lock retained engine state")
497            .push(RetainedGuard { _guard: guard });
498    }
499}
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504
505    #[derive(Debug, Clone)]
506    struct Input {
507        value: String,
508    }
509
510    #[derive(Debug, Clone)]
511    struct Processed {
512        length: usize,
513    }
514
515    #[derive(Debug, Clone)]
516    struct Final {
517        message: String,
518    }
519
520    impl From<Input> for Processed {
521        fn from(input: Input) -> Self {
522            Processed {
523                length: input.value.len(),
524            }
525        }
526    }
527
528    impl From<Processed> for Final {
529        fn from(processed: Processed) -> Self {
530            Final {
531                message: format!("Processed length: {}", processed.length),
532            }
533        }
534    }
535
536    #[test]
537    fn test_insert_and_get() {
538        let mut ctx = Context::new(Input {
539            value: "Hello".to_string(),
540        });
541
542        ctx.insert("key1", 42);
543        ctx.insert("key2", "some data".to_string());
544
545        assert_eq!(*ctx.get::<i32>("key1").unwrap(), 42);
546        assert_eq!(*ctx.get::<String>("key2").unwrap(), "some data");
547        assert!(ctx.get::<f64>("key1").is_err()); // Testing a downcast failure
548    }
549
550    #[test]
551    fn test_metadata_preserved_across_transfers() {
552        let mut ctx = Context::new(Input {
553            value: "Hello".to_string(),
554        });
555        ctx.insert_metadata("tenant", "alpha");
556
557        let (_, transferred) = ctx.transfer(Processed { length: 5 });
558        assert_eq!(
559            transferred.metadata().get("tenant").map(String::as_str),
560            Some("alpha")
561        );
562    }
563
564    #[test]
565    fn test_with_id_and_metadata_constructor() {
566        let metadata = BTreeMap::from([("tenant".to_string(), "alpha".to_string())]);
567        let ctx = Context::with_id_and_metadata(
568            Input {
569                value: "Hello".to_string(),
570            },
571            "request-123".to_string(),
572            metadata,
573        );
574
575        assert_eq!(ctx.id(), "request-123");
576        assert_eq!(
577            ctx.metadata().get("tenant").map(String::as_str),
578            Some("alpha")
579        );
580    }
581
582    #[test]
583    fn test_metadata_preserved_across_rejoin() {
584        let mut ctx = Context::new(Input {
585            value: "Hello".to_string(),
586        });
587        ctx.insert_metadata("tenant", "alpha");
588
589        let (input, empty_ctx) = ctx.into_parts();
590        let rejoined = Context::rejoin(input, empty_ctx);
591        assert_eq!(
592            rejoined.metadata().get("tenant").map(String::as_str),
593            Some("alpha")
594        );
595    }
596
597    #[test]
598    fn test_metadata_preserved_in_stream_context() {
599        let mut ctx = Context::new(Input {
600            value: "Hello".to_string(),
601        });
602        ctx.insert_metadata("tenant", "alpha");
603
604        let stream_ctx = StreamContext::from(ctx);
605        assert_eq!(
606            stream_ctx.metadata().get("tenant").map(String::as_str),
607            Some("alpha")
608        );
609    }
610
611    #[test]
612    fn test_transfer() {
613        let ctx = Context::new(Input {
614            value: "Hello".to_string(),
615        });
616
617        let (input, ctx) = ctx.transfer(Processed { length: 5 });
618
619        assert_eq!(input.value, "Hello");
620        assert_eq!(ctx.length, 5);
621    }
622
623    #[test]
624    fn test_map() {
625        let ctx = Context::new(Input {
626            value: "Hello".to_string(),
627        });
628
629        let ctx: Context<Processed> = ctx.map(|input| input.into());
630        let ctx: Context<Final> = ctx.map(|processed| processed.into());
631
632        assert_eq!(ctx.current.message, "Processed length: 5");
633    }
634
635    #[test]
636    fn test_into_context() {
637        let ctx = Context::new(Input {
638            value: "Hello".to_string(),
639        });
640
641        let ctx: Context<Processed> = ctx.into_context();
642        let ctx: Context<Final> = ctx.into_context();
643
644        assert_eq!(ctx.current.message, "Processed length: 5");
645    }
646}
dynamo_runtime/pipeline/context.rs

dynamo_runtime/pipeline/
context.rs
