awaken_runtime/extensions/background/
plugin.rs

1use std::sync::Arc;
2
3use awaken_contract::StateError;
4use awaken_contract::model::Phase;
5use awaken_contract::registry_spec::AgentSpec;
6
7use crate::plugins::{Plugin, PluginDescriptor, PluginRegistrar};
8use crate::state::{KeyScope, MutationBatch, StateKeyOptions, StateStore};
9
10use super::hook::BackgroundTaskSyncHook;
11use super::manager::BackgroundTaskManager;
12use super::state::{BackgroundTaskStateKey, BackgroundTaskViewKey};
13use super::types::BACKGROUND_TASKS_PLUGIN_ID;
14
15/// Plugin that registers the background task view state key and
16/// the persisted task metadata state key.
17pub struct BackgroundTaskPlugin {
18    manager: Arc<BackgroundTaskManager>,
19}
20
21impl BackgroundTaskPlugin {
22    pub fn new(manager: Arc<BackgroundTaskManager>) -> Self {
23        Self { manager }
24    }
25
26    /// Create the plugin and wire the store into the manager.
27    pub fn with_store(manager: Arc<BackgroundTaskManager>, store: StateStore) -> Self {
28        manager.set_store(store);
29        Self { manager }
30    }
31
32    /// Return the manager for inbox wiring.
33    pub fn manager(&self) -> &Arc<BackgroundTaskManager> {
34        &self.manager
35    }
36}
37
38impl Plugin for BackgroundTaskPlugin {
39    fn descriptor(&self) -> PluginDescriptor {
40        PluginDescriptor {
41            name: BACKGROUND_TASKS_PLUGIN_ID,
42        }
43    }
44
45    fn bind_runtime_context(
46        &self,
47        store: &StateStore,
48        owner_inbox: Option<&crate::inbox::InboxSender>,
49    ) {
50        self.manager.set_store(store.clone());
51        if let Some(inbox) = owner_inbox {
52            self.manager.set_owner_inbox(inbox.clone());
53        }
54    }
55
56    fn register(&self, registrar: &mut PluginRegistrar) -> Result<(), StateError> {
57        registrar.register_key::<BackgroundTaskViewKey>(StateKeyOptions::default())?;
58        registrar.register_key::<BackgroundTaskStateKey>(StateKeyOptions {
59            persistent: true,
60            scope: KeyScope::Thread,
61            ..StateKeyOptions::default()
62        })?;
63
64        // Sync task metadata into persisted state at run boundaries.
65        registrar.register_phase_hook(
66            BACKGROUND_TASKS_PLUGIN_ID,
67            Phase::RunStart,
68            BackgroundTaskSyncHook {
69                manager: self.manager.clone(),
70            },
71        )?;
72        registrar.register_phase_hook(
73            BACKGROUND_TASKS_PLUGIN_ID,
74            Phase::RunEnd,
75            BackgroundTaskSyncHook {
76                manager: self.manager.clone(),
77            },
78        )?;
79        // Update PendingWorkKey at step boundaries so the orchestrator
80        // can detect running tasks without knowing about this plugin.
81        registrar.register_phase_hook(
82            BACKGROUND_TASKS_PLUGIN_ID,
83            Phase::StepEnd,
84            BackgroundTaskSyncHook {
85                manager: self.manager.clone(),
86            },
87        )?;
88
89        Ok(())
90    }
91
92    fn on_activate(
93        &self,
94        _agent_spec: &AgentSpec,
95        _patch: &mut MutationBatch,
96    ) -> Result<(), StateError> {
97        Ok(())
98    }
99}
awaken_runtime/extensions/background/plugin.rs

awaken_runtime/extensions/background/
plugin.rs
