rust_web_server/a2a/
tasks.rs

1//! Stored tasks, their stream subscribers, and [`TaskHandle`], which a
2//! handler keeps to update a task from another thread.
3
4use std::collections::{HashMap, VecDeque};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::sync::{Arc, Condvar, Mutex, MutexGuard};
7
8use super::model::{self, Obj};
9use super::{Artifact, Message, Task, TaskState, TaskStatus, MAX_TASKS};
10use crate::sse::{SseEvent, SseSender};
11
12/// Streams per task; more subscribers are refused.
13pub(crate) const MAX_SUBSCRIBERS: usize = 16;
14
15pub(crate) struct Subscriber {
16    pub(crate) sender: SseSender,
17    /// The JSON-RPC id the stream's events answer.
18    pub(crate) rpc_id: String,
19}
20
21pub(crate) struct Stored {
22    pub(crate) owner: String,
23    pub(crate) task: Task,
24    pub(crate) subscribers: Vec<Subscriber>,
25    pub(crate) canceled: Arc<AtomicBool>,
26}
27
28#[derive(Default)]
29pub(crate) struct TaskStore {
30    pub(crate) tasks: HashMap<String, Stored>,
31    pub(crate) order: VecDeque<String>,
32}
33
34impl TaskStore {
35    /// Store `task` (replacing its earlier state, keeping subscribers), and
36    /// drop the oldest finished tasks past [`MAX_TASKS`].
37    pub(crate) fn put(&mut self, owner: &str, task: Task) {
38        match self.tasks.get_mut(&task.id) {
39            Some(stored) => stored.task = task,
40            None => {
41                self.order.push_back(task.id.clone());
42                self.tasks.insert(task.id.clone(), Stored { owner: owner.to_string(), task, subscribers: Vec::new(), canceled: Arc::new(AtomicBool::new(false)) });
43            }
44        }
45        while self.tasks.len() > MAX_TASKS {
46            let finished = self.order.iter().position(|id| self.tasks.get(id).is_none_or(|s| s.task.status.state.is_terminal())).unwrap_or(0);
47            match self.order.remove(finished) {
48                Some(id) => {
49                    self.tasks.remove(&id);
50                }
51                None => break,
52            }
53        }
54    }
55
56    pub(crate) fn get(&self, owner: &str, id: &str) -> Option<&Task> {
57        self.tasks.get(id).filter(|s| s.owner == owner).map(|s| &s.task)
58    }
59
60    pub(crate) fn get_mut(&mut self, owner: &str, id: &str) -> Option<&mut Stored> {
61        self.tasks.get_mut(id).filter(|s| s.owner == owner)
62    }
63}
64
65/// The store and the signal for whoever waits on a task.
66#[derive(Default)]
67pub(crate) struct Shared {
68    pub(crate) store: Mutex<TaskStore>,
69    pub(crate) changed: Condvar,
70}
71
72impl Shared {
73    pub(crate) fn lock(&self) -> MutexGuard<'_, TaskStore> {
74        self.store.lock().unwrap_or_else(|e| e.into_inner())
75    }
76}
77
78/// One stream event: a JSON-RPC response whose result is a
79/// `StreamResponse` (`task`, `message`, `statusUpdate` or `artifactUpdate`).
80pub(crate) fn event(rpc_id: &str, kind: &str, json: &str) -> SseEvent {
81    SseEvent::data(format!(r#"{{"jsonrpc":"2.0","id":{rpc_id},"result":{{"{kind}":{json}}}}}"#))
82}
83
84/// Send to every subscriber; one whose buffer is full or who left is
85/// dropped (it can subscribe again and gets the whole task first).
86pub(crate) fn publish(stored: &mut Stored, kind: &str, json: &str) {
87    stored.subscribers.retain(|s| matches!(s.sender.try_send(event(&s.rpc_id, kind, json)), Ok(true)));
88}
89
90/// Update a stored task to `state` (unless it already finished), tell its
91/// subscribers, end their streams when the task finishes or waits on the
92/// client, and wake whoever waits on it. `false` when the task is gone or
93/// already finished.
94pub(crate) fn set_state(shared: &Shared, owner: &str, id: &str, state: TaskState, message: Option<Message>) -> bool {
95    let mut store = shared.lock();
96    let Some(stored) = store.get_mut(owner, id) else { return false };
97    if stored.task.status.state.is_terminal() {
98        return false;
99    }
100    let message = message.map(|mut m| {
101        m.task_id = Some(stored.task.id.clone());
102        m.context_id = Some(stored.task.context_id.clone());
103        m
104    });
105    if let Some(m) = &message {
106        stored.task.history.push(m.clone());
107    }
108    stored.task.status = TaskStatus { state, message, timestamp: model::timestamp() };
109    let update = Obj::new().str("taskId", &stored.task.id).str("contextId", &stored.task.context_id).raw("status", &stored.task.status.to_json()).done();
110    publish(stored, "statusUpdate", &update);
111    if state.is_terminal() || state.is_interrupted() {
112        stored.subscribers.clear();
113    }
114    if state == TaskState::Canceled {
115        stored.canceled.store(true, Ordering::Relaxed);
116    }
117    drop(store);
118    shared.changed.notify_all();
119    true
120}
121
122/// A task being worked on in the background: keep it (it's `Clone` and
123/// `Send`), report progress, and finish it. Updates after the task finished
124/// (or was canceled) are ignored.
125#[derive(Clone)]
126pub struct TaskHandle {
127    pub(crate) shared: Arc<Shared>,
128    pub(crate) owner: String,
129    pub(crate) id: String,
130    pub(crate) canceled: Arc<AtomicBool>,
131}
132
133impl TaskHandle {
134    pub fn id(&self) -> &str {
135        &self.id
136    }
137
138    /// Working, with an optional progress note to the client.
139    pub fn working(&self, note: Option<&str>) {
140        set_state(&self.shared, &self.owner, &self.id, TaskState::Working, note.map(Message::agent_text));
141    }
142
143    /// Add an output; subscribers get it at once.
144    pub fn artifact(&self, artifact: Artifact) {
145        let mut store = self.shared.lock();
146        let Some(stored) = store.get_mut(&self.owner, &self.id) else { return };
147        if stored.task.status.state.is_terminal() {
148            return;
149        }
150        let update = Obj::new()
151            .str("taskId", &stored.task.id)
152            .str("contextId", &stored.task.context_id)
153            .raw("artifact", &artifact.to_json())
154            .raw("append", "false")
155            .raw("lastChunk", "true")
156            .done();
157        stored.task.artifacts.push(artifact);
158        publish(stored, "artifactUpdate", &update);
159    }
160
161    /// Done.
162    pub fn complete(&self) {
163        set_state(&self.shared, &self.owner, &self.id, TaskState::Completed, None);
164    }
165
166    /// Done with a closing message.
167    pub fn complete_with(&self, message: Message) {
168        set_state(&self.shared, &self.owner, &self.id, TaskState::Completed, Some(message));
169    }
170
171    /// Failed, saying why.
172    pub fn fail(&self, reason: &str) {
173        set_state(&self.shared, &self.owner, &self.id, TaskState::Failed, Some(Message::agent_text(reason)));
174    }
175
176    /// Waiting for the client to answer `question` (its next message
177    /// carries this task's id).
178    pub fn input_required(&self, question: &str) {
179        set_state(&self.shared, &self.owner, &self.id, TaskState::InputRequired, Some(Message::agent_text(question)));
180    }
181
182    /// The client canceled the task: stop working on it.
183    pub fn is_canceled(&self) -> bool {
184        self.canceled.load(Ordering::Relaxed)
185    }
186
187    /// The task as it is now.
188    pub fn snapshot(&self) -> Option<Task> {
189        self.shared.lock().get(&self.owner, &self.id).cloned()
190    }
191}
rust_web_server/a2a/tasks.rs

rust_web_server/a2a/
tasks.rs
