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

rust_web_server/a2a/
tasks.rs
