awaken_runtime/extensions/background/
state.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::state::StateKey;
6
7use super::types::{TaskId, TaskParentContext, TaskStatus, TaskSummary};
8
9/// Cached task view stored in the state store for prompt injection.
10#[derive(Debug, Clone, Default, Serialize, Deserialize)]
11pub struct BackgroundTaskView {
12    pub tasks: HashMap<String, TaskViewEntry>,
13}
14
15/// Lightweight view of a single background task.
16#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
17pub struct TaskViewEntry {
18    pub task_type: String,
19    pub description: String,
20    pub status: TaskStatus,
21}
22
23/// Action for the background task view state.
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub enum BackgroundTaskViewAction {
26    Replace {
27        tasks: HashMap<String, TaskViewEntry>,
28    },
29    Clear,
30}
31
32impl BackgroundTaskView {
33    pub(crate) fn reduce(&mut self, action: BackgroundTaskViewAction) {
34        match action {
35            BackgroundTaskViewAction::Replace { tasks } => {
36                self.tasks = tasks;
37            }
38            BackgroundTaskViewAction::Clear => {
39                self.tasks.clear();
40            }
41        }
42    }
43}
44
45/// State key for the cached background task view.
46pub struct BackgroundTaskViewKey;
47
48impl StateKey for BackgroundTaskViewKey {
49    const KEY: &'static str = "background_tasks";
50    type Value = BackgroundTaskView;
51    type Update = BackgroundTaskViewAction;
52
53    fn apply(value: &mut Self::Value, update: Self::Update) {
54        value.reduce(update);
55    }
56}
57
58/// Persisted metadata for a single background task.
59///
60/// Task payloads (the actual futures) are NOT persisted — only metadata
61/// (id, name, status, error message, timestamps).
62#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
63pub struct PersistedTaskMeta {
64    pub task_id: TaskId,
65    #[serde(default)]
66    pub owner_thread_id: String,
67    pub task_type: String,
68    /// Short unique name for addressing (e.g. "researcher").
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub name: Option<String>,
71    /// Human-readable task description.
72    pub description: String,
73    pub status: TaskStatus,
74    #[serde(default, skip_serializing_if = "Option::is_none")]
75    pub error: Option<String>,
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub result: Option<serde_json::Value>,
78    pub created_at_ms: u64,
79    #[serde(default, skip_serializing_if = "Option::is_none")]
80    pub completed_at_ms: Option<u64>,
81    #[serde(default)]
82    pub parent_context: TaskParentContext,
83}
84
85impl PersistedTaskMeta {
86    /// Build from a [`TaskSummary`].
87    pub fn from_summary(summary: &TaskSummary, owner_thread_id: &str) -> Self {
88        Self {
89            task_id: summary.task_id.clone(),
90            owner_thread_id: owner_thread_id.to_string(),
91            task_type: summary.task_type.clone(),
92            name: None,
93            description: summary.description.clone(),
94            status: summary.status,
95            error: summary.error.clone(),
96            result: summary.result.clone(),
97            created_at_ms: summary.created_at_ms,
98            completed_at_ms: summary.completed_at_ms,
99            parent_context: summary.parent_context.clone(),
100        }
101    }
102}
103
104/// Persisted state for all background tasks.
105#[derive(Debug, Clone, Default, Serialize, Deserialize)]
106pub struct BackgroundTaskStateSnapshot {
107    pub tasks: HashMap<TaskId, PersistedTaskMeta>,
108}
109
110/// Actions applied to the persisted background task state.
111#[derive(Debug, Clone, Serialize, Deserialize)]
112pub enum BackgroundTaskStateAction {
113    /// Upsert a single task's metadata.
114    Upsert(Box<PersistedTaskMeta>),
115    /// Replace the entire task map (used on restore/sync).
116    ReplaceAll {
117        tasks: HashMap<TaskId, PersistedTaskMeta>,
118    },
119}
120
121impl BackgroundTaskStateSnapshot {
122    pub(crate) fn reduce(&mut self, action: BackgroundTaskStateAction) {
123        match action {
124            BackgroundTaskStateAction::Upsert(meta) => {
125                self.tasks.insert(meta.task_id.clone(), *meta);
126            }
127            BackgroundTaskStateAction::ReplaceAll { tasks } => {
128                self.tasks = tasks;
129            }
130        }
131    }
132}
133
134/// State key for persisted background task metadata.
135///
136/// Scoped to `Thread` so it survives across runs. On task completion or
137/// failure the manager writes a state update; on resume, the plugin
138/// restores known task metadata from this key.
139pub struct BackgroundTaskStateKey;
140
141impl StateKey for BackgroundTaskStateKey {
142    const KEY: &'static str = "background_task_state";
143    type Value = BackgroundTaskStateSnapshot;
144    type Update = BackgroundTaskStateAction;
145
146    fn apply(value: &mut Self::Value, update: Self::Update) {
147        value.reduce(update);
148    }
149}
awaken_runtime/extensions/background/state.rs

awaken_runtime/extensions/background/
state.rs
