1use crate::types::RequestId;
7use std::collections::HashMap;
8
9#[derive(Debug, Clone, PartialEq, Eq, Hash)]
10pub enum RequestState {
11 Init,
12 Running,
13 Streaming,
14 Completed,
15 Cancelled,
16 Error,
17}
18
19pub struct RequestTable {
21 requests: HashMap<RequestId, RequestState>,
22}
23
24impl Default for RequestTable {
25 fn default() -> Self {
26 Self::new()
27 }
28}
29
30impl RequestTable {
31 #[must_use]
32 pub fn new() -> Self {
33 Self {
34 requests: HashMap::new(),
35 }
36 }
37
38 pub fn start(&mut self, id: RequestId) {
39 self.requests.insert(id, RequestState::Running);
40 }
41
42 pub fn mark_streaming(&mut self, id: RequestId) {
43 if let Some(state) = self.requests.get_mut(&id)
44 && *state == RequestState::Running
45 {
46 *state = RequestState::Streaming;
47 }
48 }
49
50 pub fn complete(&mut self, id: RequestId) {
51 self.requests.insert(id, RequestState::Completed);
52 }
53
54 pub fn cancel(&mut self, id: RequestId) -> bool {
55 match self.requests.get_mut(&id) {
56 Some(RequestState::Running | RequestState::Streaming) => {
57 self.requests.insert(id, RequestState::Cancelled);
58 true
59 }
60 _ => false,
61 }
62 }
63
64 #[must_use]
65 pub fn is_cancelled(&self, id: RequestId) -> bool {
66 matches!(self.requests.get(&id), Some(RequestState::Cancelled))
67 }
68
69 pub fn cleanup(&mut self, id: RequestId) {
70 self.requests.remove(&id);
71 }
72}