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;
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
334impl AsyncEngineContextProvider for StreamContext {
335    fn context(&self) -> Arc<dyn AsyncEngineContext> {
336        self.controller.clone()
337    }
338}
339
340impl<T: Send + Sync + 'static> From<Context<T>> for StreamContext {
341    fn from(value: Context<T>) -> Self {
342        StreamContext::new(value.controller, value.registry, value.metadata)
343    }
344}
345
346// TODO - refactor here - this came from the dynamo.llm-async-engine crate
347
348use tokio::sync::watch::{Receiver, Sender, channel};
349
350#[derive(Debug, Eq, PartialEq)]
351enum State {
352    Live,
353    Stopped,
354    Killed,
355}
356
357/// A context implementation with cancellation propagation.
358#[derive(Debug)]
359pub struct Controller {
360    id: String,
361    tx: Sender<State>,
362    rx: Receiver<State>,
363    child_context: Mutex<Vec<Arc<dyn AsyncEngineContext>>>,
364}
365
366impl Controller {
367    pub fn new(id: String) -> Self {
368        let (tx, rx) = channel(State::Live);
369        Self {
370            id,
371            tx,
372            rx,
373            child_context: Mutex::new(Vec::new()),
374        }
375    }
376
377    pub fn id(&self) -> &str {
378        &self.id
379    }
380}
381
382impl Default for Controller {
383    fn default() -> Self {
384        Self::new(uuid::Uuid::new_v4().to_string())
385    }
386}
387
388impl AsyncEngineController for Controller {}
389
390#[async_trait]
391impl AsyncEngineContext for Controller {
392    fn id(&self) -> &str {
393        &self.id
394    }
395
396    fn is_stopped(&self) -> bool {
397        *self.rx.borrow() != State::Live
398    }
399
400    fn is_killed(&self) -> bool {
401        *self.rx.borrow() == State::Killed
402    }
403
404    async fn stopped(&self) {
405        let mut rx = self.rx.clone();
406        loop {
407            if *rx.borrow_and_update() != State::Live || rx.changed().await.is_err() {
408                return;
409            }
410        }
411    }
412
413    async fn killed(&self) {
414        let mut rx = self.rx.clone();
415        loop {
416            if *rx.borrow_and_update() == State::Killed || rx.changed().await.is_err() {
417                return;
418            }
419        }
420    }
421
422    fn stop_generating(&self) {
423        // Clone child Arcs to avoid deadlock if parent is accidentally linked under child
424        let children = self
425            .child_context
426            .lock()
427            .expect("Failed to lock child context")
428            .iter()
429            .cloned()
430            .collect::<Vec<_>>();
431        for child in children {
432            child.stop_generating();
433        }
434
435        let _ = self.tx.send(State::Stopped);
436    }
437
438    fn stop(&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();
449        }
450
451        let _ = self.tx.send(State::Stopped);
452    }
453
454    fn kill(&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.kill();
465        }
466
467        let _ = self.tx.send(State::Killed);
468    }
469
470    fn link_child(&self, child: Arc<dyn AsyncEngineContext>) {
471        self.child_context
472            .lock()
473            .expect("Failed to lock child context")
474            .push(child);
475    }
476}
477
478#[cfg(test)]
479mod tests {
480    use super::*;
481
482    #[derive(Debug, Clone)]
483    struct Input {
484        value: String,
485    }
486
487    #[derive(Debug, Clone)]
488    struct Processed {
489        length: usize,
490    }
491
492    #[derive(Debug, Clone)]
493    struct Final {
494        message: String,
495    }
496
497    impl From<Input> for Processed {
498        fn from(input: Input) -> Self {
499            Processed {
500                length: input.value.len(),
501            }
502        }
503    }
504
505    impl From<Processed> for Final {
506        fn from(processed: Processed) -> Self {
507            Final {
508                message: format!("Processed length: {}", processed.length),
509            }
510        }
511    }
512
513    #[test]
514    fn test_insert_and_get() {
515        let mut ctx = Context::new(Input {
516            value: "Hello".to_string(),
517        });
518
519        ctx.insert("key1", 42);
520        ctx.insert("key2", "some data".to_string());
521
522        assert_eq!(*ctx.get::<i32>("key1").unwrap(), 42);
523        assert_eq!(*ctx.get::<String>("key2").unwrap(), "some data");
524        assert!(ctx.get::<f64>("key1").is_err()); // Testing a downcast failure
525    }
526
527    #[test]
528    fn test_metadata_preserved_across_transfers() {
529        let mut ctx = Context::new(Input {
530            value: "Hello".to_string(),
531        });
532        ctx.insert_metadata("tenant", "alpha");
533
534        let (_, transferred) = ctx.transfer(Processed { length: 5 });
535        assert_eq!(
536            transferred.metadata().get("tenant").map(String::as_str),
537            Some("alpha")
538        );
539    }
540
541    #[test]
542    fn test_with_id_and_metadata_constructor() {
543        let metadata = BTreeMap::from([("tenant".to_string(), "alpha".to_string())]);
544        let ctx = Context::with_id_and_metadata(
545            Input {
546                value: "Hello".to_string(),
547            },
548            "request-123".to_string(),
549            metadata,
550        );
551
552        assert_eq!(ctx.id(), "request-123");
553        assert_eq!(
554            ctx.metadata().get("tenant").map(String::as_str),
555            Some("alpha")
556        );
557    }
558
559    #[test]
560    fn test_metadata_preserved_across_rejoin() {
561        let mut ctx = Context::new(Input {
562            value: "Hello".to_string(),
563        });
564        ctx.insert_metadata("tenant", "alpha");
565
566        let (input, empty_ctx) = ctx.into_parts();
567        let rejoined = Context::rejoin(input, empty_ctx);
568        assert_eq!(
569            rejoined.metadata().get("tenant").map(String::as_str),
570            Some("alpha")
571        );
572    }
573
574    #[test]
575    fn test_metadata_preserved_in_stream_context() {
576        let mut ctx = Context::new(Input {
577            value: "Hello".to_string(),
578        });
579        ctx.insert_metadata("tenant", "alpha");
580
581        let stream_ctx = StreamContext::from(ctx);
582        assert_eq!(
583            stream_ctx.metadata().get("tenant").map(String::as_str),
584            Some("alpha")
585        );
586    }
587
588    #[test]
589    fn test_transfer() {
590        let ctx = Context::new(Input {
591            value: "Hello".to_string(),
592        });
593
594        let (input, ctx) = ctx.transfer(Processed { length: 5 });
595
596        assert_eq!(input.value, "Hello");
597        assert_eq!(ctx.length, 5);
598    }
599
600    #[test]
601    fn test_map() {
602        let ctx = Context::new(Input {
603            value: "Hello".to_string(),
604        });
605
606        let ctx: Context<Processed> = ctx.map(|input| input.into());
607        let ctx: Context<Final> = ctx.map(|processed| processed.into());
608
609        assert_eq!(ctx.current.message, "Processed length: 5");
610    }
611
612    #[test]
613    fn test_into_context() {
614        let ctx = Context::new(Input {
615            value: "Hello".to_string(),
616        });
617
618        let ctx: Context<Processed> = ctx.into_context();
619        let ctx: Context<Final> = ctx.into_context();
620
621        assert_eq!(ctx.current.message, "Processed length: 5");
622    }
623}
dynamo_runtime/pipeline/context.rs

dynamo_runtime/pipeline/
context.rs
