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

rws_proxy/a2a/
tasks.rs
