Skip to main content

arbiter_session/
middleware.rs

1//! Session middleware utilities for proxy integration.
2//!
3//! Provides session ID parsing from the `X-Arbiter-Session` header and
4//! error-to-status-code mapping. Session constraint validation is performed
5//! directly by the handler stages in `crates/arbiter/src/stages/`.
6
7use uuid::Uuid;
8
9use crate::error::SessionError;
10
11/// HTTP status code that should be returned for each session error type.
12pub fn status_code_for_error(err: &SessionError) -> u16 {
13    match err {
14        SessionError::NotFound(_) => 404,
15        SessionError::Expired(_) => 408,
16        SessionError::BudgetExceeded { .. } => 429,
17        SessionError::ToolNotAuthorized { .. } => 403,
18        SessionError::AlreadyClosed(_) => 410,
19        SessionError::RateLimited { .. } => 429,
20        SessionError::TooManySessions { .. } => 429,
21    }
22}
23
24/// Extract a session ID from the `X-Arbiter-Session` header value.
25///
26/// Returns `None` if the header is missing or not a valid UUID.
27pub fn parse_session_header(header_value: &str) -> Option<Uuid> {
28    Uuid::parse_str(header_value.trim()).ok()
29}
30
31#[cfg(test)]
32mod tests {
33    use super::*;
34
35    #[test]
36    fn parse_session_header_valid() {
37        let id = Uuid::new_v4();
38        assert_eq!(parse_session_header(&id.to_string()), Some(id));
39    }
40
41    #[test]
42    fn parse_session_header_invalid() {
43        assert_eq!(parse_session_header("not-a-uuid"), None);
44        assert_eq!(parse_session_header(""), None);
45    }
46
47    /// Test that each SessionError variant maps to the correct HTTP status code.
48    #[test]
49    fn status_code_for_all_error_variants() {
50        let id = Uuid::new_v4();
51
52        assert_eq!(status_code_for_error(&SessionError::NotFound(id)), 404);
53        assert_eq!(status_code_for_error(&SessionError::Expired(id)), 408);
54        assert_eq!(
55            status_code_for_error(&SessionError::BudgetExceeded {
56                session_id: id,
57                limit: 10,
58                used: 10,
59            }),
60            429
61        );
62        assert_eq!(
63            status_code_for_error(&SessionError::ToolNotAuthorized {
64                session_id: id,
65                tool: "delete".into(),
66            }),
67            403
68        );
69        assert_eq!(status_code_for_error(&SessionError::AlreadyClosed(id)), 410);
70        assert_eq!(
71            status_code_for_error(&SessionError::RateLimited {
72                session_id: id,
73                limit_per_minute: 5,
74            }),
75            429
76        );
77        assert_eq!(
78            status_code_for_error(&SessionError::TooManySessions {
79                agent_id: id.to_string(),
80                max: 10,
81                current: 10,
82            }),
83            429
84        );
85    }
86
87    /// Whitespace-padded UUIDs should be accepted (the trim() in parse_session_header).
88    #[test]
89    fn parse_session_header_with_whitespace() {
90        let id = Uuid::new_v4();
91        let padded = format!("  {}  ", id);
92        assert_eq!(
93            parse_session_header(&padded),
94            Some(id),
95            "whitespace-padded UUID should be parsed correctly"
96        );
97    }
98}